/*
 * Copyright (C) 2024 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package android.security.cts.CVE_2024_43769;

import static android.Manifest.permission.BYPASS_ROLE_QUALIFICATION;
import static android.Manifest.permission.MANAGE_PROFILE_AND_DEVICE_OWNERS;
import static android.Manifest.permission.MANAGE_ROLE_HOLDERS;

import static androidx.test.core.app.ApplicationProvider.getApplicationContext;

import static com.android.compatibility.common.util.SystemUtil.runShellCommand;
import static com.android.compatibility.common.util.SystemUtil.runWithShellPermissionIdentity;

import static com.google.common.truth.Truth.assertWithMessage;
import static com.google.common.truth.TruthJUnit.assume;

import android.app.admin.DevicePolicyManager;
import android.app.role.RoleManager;
import android.content.Context;
import android.os.UserHandle;

import androidx.test.runner.AndroidJUnit4;

import org.junit.Test;
import org.junit.runner.RunWith;

import java.util.List;
import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

@RunWith(AndroidJUnit4.class)
public class DeviceTest {
    final Semaphore mIsRoleGranted = new Semaphore(0);

    @Test
    public void testPocCVE_2024_43769() {
        try {
            final Context context = getApplicationContext();
            final RoleManager roleManager = context.getSystemService(RoleManager.class);

            // Check if any managed user is present on device
            final List<UserHandle> managedProfiles =
                    runWithShellPermissionIdentity(
                            () -> {
                                return context.getSystemService(DevicePolicyManager.class)
                                        .getPolicyManagedProfiles(
                                                UserHandle.of(UserHandle.myUserId()));
                            },
                            MANAGE_PROFILE_AND_DEVICE_OWNERS);
            assume().withMessage("Device should not have managed user !!")
                    .that(managedProfiles.isEmpty())
                    .isTrue();

            // Bypass role qualification to add 'DEVICE_POLICY_MANAGEMENT' role and assign
            // 'DEVICE_POLICY_MANAGEMENT' role holder to the helper application to make it a
            // 'Device Admin'
            String packageName = "android.security.cts.CVE_2024_43769_helper";
            try (AutoCloseable withBypassingRoleQualificationTrue =
                            withBypassingRoleQualificationTrue(roleManager);
                    AutoCloseable withRoleHolder =
                            withRoleHolder(
                                    roleManager,
                                    context,
                                    RoleManager.ROLE_DEVICE_POLICY_MANAGEMENT,
                                    packageName)) {
                // Check whether 'DEVICE_POLICY_MANAGEMENT' is granted
                assume().withMessage("Failed to grant 'DEVICE_POLICY_MANAGEMENT' role!!")
                        .that(mIsRoleGranted.tryAcquire(10, TimeUnit.SECONDS))
                        .isTrue();

                // Uninstall helper application
                String result =
                        runShellCommand("pm uninstall android.security.cts.CVE_2024_43769_helper");

                // Create a regex pattern to check whether 'DELETE_FAILED_DEVICE_POLICY_MANAGER' is
                // present in the result
                Pattern pattern =
                        Pattern.compile(
                                "DELETE_FAILED_DEVICE_POLICY_MANAGER", Pattern.CASE_INSENSITIVE);
                Matcher matcher = pattern.matcher(result);

                // Without fix, the installed helper application will not get uninstalled
                assertWithMessage(
                                "Device is vulnerable to b/360807442 !! CloudDPC can't be"
                                        + " uninstalled on unamanaged devices")
                        .that(matcher.find())
                        .isFalse();
            }
        } catch (Exception e) {
            assume().that(e).isNull();
        }
    }

    private AutoCloseable withBypassingRoleQualificationTrue(RoleManager roleManager)
            throws Exception {
        final boolean isBypassingRoleQualificationEnabled =
                runWithShellPermissionIdentity(
                        () -> {
                            return roleManager.isBypassingRoleQualification();
                        },
                        MANAGE_ROLE_HOLDERS);

        if (isBypassingRoleQualificationEnabled) {
            return () -> {};
        }

        // Set bypass role qualification to true
        runWithShellPermissionIdentity(
                () -> {
                    roleManager.setBypassingRoleQualification(true /* bypassRoleQualification */);
                },
                BYPASS_ROLE_QUALIFICATION);

        // Set bypass role qualification to false on return
        return () ->
                runWithShellPermissionIdentity(
                        () -> {
                            roleManager.setBypassingRoleQualification(
                                    false /* bypassRoleQualification */);
                        },
                        BYPASS_ROLE_QUALIFICATION);
    }

    private AutoCloseable withRoleHolder(
            RoleManager roleManager, Context context, String roleName, String packageName)
            throws Exception {
        // Assign 'DEVICE_POLICY_MANAGEMENT' role to the package
        runWithShellPermissionIdentity(
                () -> {
                    roleManager.addRoleHolderAsUser(
                            roleName,
                            packageName,
                            0 /* flags */,
                            UserHandle.SYSTEM /* user */,
                            context.getMainExecutor() /* Executor */,
                            successful -> {
                                if (successful) {
                                    mIsRoleGranted.release();
                                }
                            } /* callback */);
                },
                MANAGE_ROLE_HOLDERS);

        // Remove the 'DEVICE_POLICY_MANAGEMENT' role on return
        return () ->
                runWithShellPermissionIdentity(
                        () -> {
                            roleManager.removeRoleHolderAsUser(
                                    roleName,
                                    packageName,
                                    0 /* flags */,
                                    UserHandle.SYSTEM /* user */,
                                    context.getMainExecutor() /* Executor */,
                                    successful -> {} /* callback */);
                        },
                        MANAGE_ROLE_HOLDERS);
    }
}