Retain FRP data when FR is invoked in Settings

* If the device is an organization-owned managed
  profile device and a FRP policy is set, the
  factory reset protection data is no longer
  erased from factory reset in Settings.

Bug: 148847767
Test: manual testing
Change-Id: Iebaf2446f23626b24db37f5173dc77f9dee25ba9
diff --git a/src/com/android/settings/MasterClearConfirm.java b/src/com/android/settings/MasterClearConfirm.java
index 679f18f..cac18f7 100644
--- a/src/com/android/settings/MasterClearConfirm.java
+++ b/src/com/android/settings/MasterClearConfirm.java
@@ -22,6 +22,8 @@
 import android.app.ActionBar;
 import android.app.Activity;
 import android.app.ProgressDialog;
+import android.app.admin.DevicePolicyManager;
+import android.app.admin.FactoryResetProtectionPolicy;
 import android.app.settings.SettingsEnums;
 import android.content.Context;
 import android.content.Intent;
@@ -83,14 +85,9 @@
 
             final PersistentDataBlockManager pdbManager = (PersistentDataBlockManager)
                     getActivity().getSystemService(Context.PERSISTENT_DATA_BLOCK_SERVICE);
-            final OemLockManager oemLockManager = (OemLockManager)
-                    getActivity().getSystemService(Context.OEM_LOCK_SERVICE);
 
-            if (pdbManager != null && !oemLockManager.isOemUnlockAllowed() &&
-                    WizardManagerHelper.isDeviceProvisioned(getActivity())) {
-                // if OEM unlock is allowed, the persistent data block will be wiped during FR
-                // process. If disabled, it will be wiped here, unless the device is still being
-                // provisioned, in which case the persistent data block will be preserved.
+            if (shouldWipePersistentDataBlock(pdbManager)) {
+
                 new AsyncTask<Void, Void, Void>() {
                     int mOldOrientation;
                     ProgressDialog mProgressDialog;
@@ -140,6 +137,44 @@
         }
     };
 
+    @VisibleForTesting
+    boolean shouldWipePersistentDataBlock(PersistentDataBlockManager pdbManager) {
+        if (pdbManager == null) {
+            return false;
+        }
+        // The persistent data block will persist if the device is still being provisioned.
+        if (isDeviceStillBeingProvisioned()) {
+            return false;
+        }
+        // If OEM unlock is allowed, the persistent data block will be wiped during FR
+        // process. If disabled, it will be wiped here instead.
+        if (isOemUnlockedAllowed()) {
+            return false;
+        }
+        // Do not erase the factory reset protection data (from Settings) if the
+        // device is an organization-owned managed profile device and a factory
+        // reset protection policy has been set.
+        final DevicePolicyManager dpm = (DevicePolicyManager) getActivity()
+                .getSystemService(Context.DEVICE_POLICY_SERVICE);
+        FactoryResetProtectionPolicy frpPolicy = dpm.getFactoryResetProtectionPolicy(null);
+        if (dpm.isOrganizationOwnedDeviceWithManagedProfile() && frpPolicy != null
+                && frpPolicy.isNotEmpty()) {
+            return false;
+        }
+        return true;
+    }
+
+    @VisibleForTesting
+    boolean isOemUnlockedAllowed() {
+        return ((OemLockManager) getActivity().getSystemService(
+                Context.OEM_LOCK_SERVICE)).isOemUnlockAllowed();
+    }
+
+    @VisibleForTesting
+    boolean isDeviceStillBeingProvisioned() {
+        return !WizardManagerHelper.isDeviceProvisioned(getActivity());
+    }
+
     private void doMasterClear() {
         Intent intent = new Intent(Intent.ACTION_FACTORY_RESET);
         intent.setPackage("android");
diff --git a/tests/robotests/src/com/android/settings/MasterClearConfirmTest.java b/tests/robotests/src/com/android/settings/MasterClearConfirmTest.java
index 822eb3c..1222bdd 100644
--- a/tests/robotests/src/com/android/settings/MasterClearConfirmTest.java
+++ b/tests/robotests/src/com/android/settings/MasterClearConfirmTest.java
@@ -18,23 +18,50 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
-import android.app.Activity;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.when;
+
+import android.app.admin.DevicePolicyManager;
+import android.app.admin.FactoryResetProtectionPolicy;
+import android.content.Context;
+import android.service.persistentdata.PersistentDataBlockManager;
 import android.view.LayoutInflater;
 import android.widget.TextView;
 
+import androidx.fragment.app.FragmentActivity;
+
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
 import org.robolectric.Robolectric;
 import org.robolectric.RobolectricTestRunner;
 
+import java.util.ArrayList;
+
 @RunWith(RobolectricTestRunner.class)
 public class MasterClearConfirmTest {
-    private Activity mActivity;
+
+    private FragmentActivity mActivity;
+
+    @Mock
+    private FragmentActivity mMockActivity;
+
+    @Mock
+    private DevicePolicyManager mDevicePolicyManager;
+
+    @Mock
+    private PersistentDataBlockManager mPersistentDataBlockManager;
+
+    private MasterClearConfirm mMasterClearConfirm;
 
     @Before
     public void setUp() {
-        mActivity = Robolectric.setupActivity(Activity.class);
+        MockitoAnnotations.initMocks(this);
+        mActivity = Robolectric.setupActivity(FragmentActivity.class);
+        mMasterClearConfirm = spy(new MasterClearConfirm());
     }
 
     @Test
@@ -64,4 +91,63 @@
                 .findViewById(R.id.sud_layout_description)).getText())
                 .isEqualTo(mActivity.getString(R.string.master_clear_final_desc));
     }
+
+    @Test
+    public void shouldWipePersistentDataBlock_noPersistentDataBlockManager_shouldReturnFalse() {
+        assertThat(mMasterClearConfirm.shouldWipePersistentDataBlock(null)).isFalse();
+    }
+
+    @Test
+    public void shouldWipePersistentDataBlock_deviceIsStillBeingProvisioned_shouldReturnFalse() {
+        doReturn(true).when(mMasterClearConfirm).isDeviceStillBeingProvisioned();
+
+        assertThat(mMasterClearConfirm.shouldWipePersistentDataBlock(
+                mPersistentDataBlockManager)).isFalse();
+    }
+
+    @Test
+    public void shouldWipePersistentDataBlock_oemUnlockAllowed_shouldReturnFalse() {
+        doReturn(false).when(mMasterClearConfirm).isDeviceStillBeingProvisioned();
+        doReturn(true).when(mMasterClearConfirm).isOemUnlockedAllowed();
+
+        assertThat(mMasterClearConfirm.shouldWipePersistentDataBlock(
+                mPersistentDataBlockManager)).isFalse();
+    }
+
+    @Test
+    public void shouldWipePersistentDataBlock_hasFactoryResetProtectionPolicy_shouldReturnFalse() {
+        when(mMasterClearConfirm.getActivity()).thenReturn(mMockActivity);
+
+        doReturn(false).when(mMasterClearConfirm).isDeviceStillBeingProvisioned();
+        doReturn(false).when(mMasterClearConfirm).isOemUnlockedAllowed();
+        ArrayList<String> accounts = new ArrayList<>();
+        accounts.add("test");
+        FactoryResetProtectionPolicy frp = new FactoryResetProtectionPolicy.Builder()
+                .setFactoryResetProtectionAccounts(accounts)
+                .setFactoryResetProtectionEnabled(true)
+                .build();
+        when(mMockActivity.getSystemService(Context.DEVICE_POLICY_SERVICE))
+                .thenReturn(mDevicePolicyManager);
+        when(mDevicePolicyManager.getFactoryResetProtectionPolicy(null)).thenReturn(frp);
+        when(mDevicePolicyManager.isOrganizationOwnedDeviceWithManagedProfile()).thenReturn(true);
+
+        assertThat(mMasterClearConfirm.shouldWipePersistentDataBlock(
+                mPersistentDataBlockManager)).isFalse();
+    }
+
+    @Test
+    public void shouldWipePersistentDataBlock_isNotOrganizationOwnedDevice_shouldReturnTrue() {
+        when(mMasterClearConfirm.getActivity()).thenReturn(mMockActivity);
+
+        doReturn(false).when(mMasterClearConfirm).isDeviceStillBeingProvisioned();
+        doReturn(false).when(mMasterClearConfirm).isOemUnlockedAllowed();
+
+        when(mMockActivity.getSystemService(Context.DEVICE_POLICY_SERVICE))
+                .thenReturn(mDevicePolicyManager);
+        when(mDevicePolicyManager.getFactoryResetProtectionPolicy(null)).thenReturn(null);
+        when(mDevicePolicyManager.isOrganizationOwnedDeviceWithManagedProfile()).thenReturn(false);
+
+        assertThat(mMasterClearConfirm.shouldWipePersistentDataBlock(
+                mPersistentDataBlockManager)).isTrue();
+    }
 }