DPM: Implement Device ID attestation

Enable requesting inclusion of device identifiers in the attestation
record issued for keys generated by generateKeyPair.
This is done by passing an array of flags with values indicating which
identifiers should be included.
Since the attestation record will include sensitive identifiers, it can
only be requested by the DPC in Device Owner mode or by the Delegated
Cert Installer in Device Owner mode.

Design note:
DevicePolicyManager defines its own set of constants for the different
identifier types (ID_TYPE_*) and prior to calling
DevicePolicyManagerService it translates them to the values defined by
AttestationUtils (which is not a public class).
The reason is to allow re-use of code in AttestationUtils for preparing
the attestation arguments.
In theory, these constants could be moved from AttestationUtils to
DevicePolicyManager, however that would create a dependency on DPM from
Keystore, which logically does not make sense as Keystore is independent
of the DPM (and in a lower level of the system, conceptually).

Bug: 63388672
Test: cts-tradefed run commandAndExit cts-dev -a armeabi-v7a -m CtsDevicePolicyManagerTestCases -t com.android.cts.devicepolicy.DeviceOwnerTest#testKeyManagement; runtest frameworks-services -c com.android.server.devicepolicy.DevicePolicyManagerTest#testTranslationOfIdAttestationFlag
Change-Id: Ifb42e8e813fa812a08203b4a81d15b1f91152354
diff --git a/api/current.txt b/api/current.txt
index 8925846..641b6e4 100644
--- a/api/current.txt
+++ b/api/current.txt
@@ -6338,7 +6338,7 @@
     method public android.os.UserHandle createAndManageUser(android.content.ComponentName, java.lang.String, android.content.ComponentName, android.os.PersistableBundle, int);
     method public void enableSystemApp(android.content.ComponentName, java.lang.String);
     method public int enableSystemApp(android.content.ComponentName, android.content.Intent);
-    method public android.security.AttestedKeyPair generateKeyPair(android.content.ComponentName, java.lang.String, android.security.keystore.KeyGenParameterSpec);
+    method public android.security.AttestedKeyPair generateKeyPair(android.content.ComponentName, java.lang.String, android.security.keystore.KeyGenParameterSpec, int);
     method public java.lang.String[] getAccountTypesWithManagementDisabled();
     method public java.util.List<android.content.ComponentName> getActiveAdmins();
     method public java.util.Set<java.lang.String> getAffiliationIds(android.content.ComponentName);
@@ -6572,6 +6572,10 @@
     field public static final int FLAG_EVICT_CREDENTIAL_ENCRYPTION_KEY = 1; // 0x1
     field public static final int FLAG_MANAGED_CAN_ACCESS_PARENT = 2; // 0x2
     field public static final int FLAG_PARENT_CAN_ACCESS_MANAGED = 1; // 0x1
+    field public static final int ID_TYPE_BASE_INFO = 1; // 0x1
+    field public static final int ID_TYPE_IMEI = 4; // 0x4
+    field public static final int ID_TYPE_MEID = 8; // 0x8
+    field public static final int ID_TYPE_SERIAL = 2; // 0x2
     field public static final int KEYGUARD_DISABLE_FEATURES_ALL = 2147483647; // 0x7fffffff
     field public static final int KEYGUARD_DISABLE_FEATURES_NONE = 0; // 0x0
     field public static final int KEYGUARD_DISABLE_FINGERPRINT = 32; // 0x20
diff --git a/core/java/android/app/admin/DevicePolicyManager.java b/core/java/android/app/admin/DevicePolicyManager.java
index 7e80ac7..0b74741 100644
--- a/core/java/android/app/admin/DevicePolicyManager.java
+++ b/core/java/android/app/admin/DevicePolicyManager.java
@@ -1668,6 +1668,46 @@
     public static final String ACTION_DEVICE_ADMIN_SERVICE
             = "android.app.action.DEVICE_ADMIN_SERVICE";
 
+    /** @hide */
+    @Retention(RetentionPolicy.SOURCE)
+    @IntDef(flag = true, prefix = {"ID_TYPE_"}, value = {
+        ID_TYPE_BASE_INFO,
+        ID_TYPE_SERIAL,
+        ID_TYPE_IMEI,
+        ID_TYPE_MEID
+    })
+    public @interface AttestationIdType {}
+
+    /**
+     * Specifies that the device should attest its manufacturer details. For use with
+     * {@link #generateKeyPair}.
+     *
+     * @see #generateKeyPair
+     */
+    public static final int ID_TYPE_BASE_INFO = 1;
+
+    /**
+     * Specifies that the device should attest its serial number. For use with
+     * {@link #generateKeyPair}.
+     *
+     * @see #generateKeyPair
+     */
+    public static final int ID_TYPE_SERIAL = 2;
+
+    /**
+     * Specifies that the device should attest its IMEI. For use with {@link #generateKeyPair}.
+     *
+     * @see #generateKeyPair
+     */
+    public static final int ID_TYPE_IMEI = 4;
+
+    /**
+     * Specifies that the device should attest its MEID. For use with {@link #generateKeyPair}.
+     *
+     * @see #generateKeyPair
+     */
+    public static final int ID_TYPE_MEID = 8;
+
     /**
      * Return true if the given administrator component is currently active (enabled) in the system.
      *
@@ -4106,22 +4146,46 @@
      * @param algorithm The key generation algorithm, see {@link java.security.KeyPairGenerator}.
      * @param keySpec Specification of the key to generate, see
      * {@link java.security.KeyPairGenerator}.
+     * @param idAttestationFlags A bitmask of all the identifiers that should be included in the
+     *        attestation record ({@code ID_TYPE_BASE_INFO}, {@code ID_TYPE_SERIAL},
+     *        {@code ID_TYPE_IMEI} and {@code ID_TYPE_MEID}), or {@code 0} if no device
+     *        identification is required in the attestation record.
+     *        Device owner, profile owner and their delegated certificate installer can use
+     *        {@link #ID_TYPE_BASE_INFO} to request inclusion of the general device information
+     *        including manufacturer, model, brand, device and product in the attestation record.
+     *        Only device owner and their delegated certificate installer can use
+     *        {@link #ID_TYPE_SERIAL}, {@link #ID_TYPE_IMEI} and {@link #ID_TYPE_MEID} to request
+     *        unique device identifiers to be attested.
+     *        <p>
+     *        If any of {@link #ID_TYPE_SERIAL}, {@link #ID_TYPE_IMEI} and {@link #ID_TYPE_MEID}
+     *        is set, it is implicitly assumed that {@link #ID_TYPE_BASE_INFO} is also set.
+     *        <p>
+     *        If any flag is specified, then an attestation challenge must be included in the
+     *        {@code keySpec}.
      * @return A non-null {@code AttestedKeyPair} if the key generation succeeded, null otherwise.
      * @throws SecurityException if {@code admin} is not {@code null} and not a device or profile
-     *         owner.
-     * @throws IllegalArgumentException if the alias in {@code keySpec} is empty, or if the
+     *         owner. If Device ID attestation is requested (using {@link #ID_TYPE_SERIAL},
+     *         {@link #ID_TYPE_IMEI} or {@link #ID_TYPE_MEID}), the caller must be the Device Owner
+     *         or the Certificate Installer delegate.
+     * @throws IllegalArgumentException if the alias in {@code keySpec} is empty, if the
      *         algorithm specification in {@code keySpec} is not {@code RSAKeyGenParameterSpec}
-     *         or {@code ECGenParameterSpec}.
+     *         or {@code ECGenParameterSpec}, or if Device ID attestation was requested but the
+     *         {@code keySpec} does not contain an attestation challenge.
+     * @see KeyGenParameterSpec.Builder#setAttestationChallenge(byte[])
      */
     public AttestedKeyPair generateKeyPair(@Nullable ComponentName admin,
-            @NonNull String algorithm, @NonNull KeyGenParameterSpec keySpec) {
+            @NonNull String algorithm, @NonNull KeyGenParameterSpec keySpec,
+            @AttestationIdType int idAttestationFlags) {
         throwIfParentInstance("generateKeyPair");
         try {
             final ParcelableKeyGenParameterSpec parcelableSpec =
                     new ParcelableKeyGenParameterSpec(keySpec);
             KeymasterCertificateChain attestationChain = new KeymasterCertificateChain();
+
+            // Translate ID attestation flags to values used by AttestationUtils
             final boolean success = mService.generateKeyPair(
-                    admin, mContext.getPackageName(), algorithm, parcelableSpec, attestationChain);
+                    admin, mContext.getPackageName(), algorithm, parcelableSpec,
+                    idAttestationFlags, attestationChain);
             if (!success) {
                 Log.e(TAG, "Error generating key via DevicePolicyManagerService.");
                 return null;
diff --git a/core/java/android/app/admin/IDevicePolicyManager.aidl b/core/java/android/app/admin/IDevicePolicyManager.aidl
index 7cf19ee..5916a62 100644
--- a/core/java/android/app/admin/IDevicePolicyManager.aidl
+++ b/core/java/android/app/admin/IDevicePolicyManager.aidl
@@ -173,7 +173,7 @@
     boolean removeKeyPair(in ComponentName who, in String callerPackage, String alias);
     boolean generateKeyPair(in ComponentName who, in String callerPackage, in String algorithm,
             in ParcelableKeyGenParameterSpec keySpec,
-            out KeymasterCertificateChain attestationChain);
+            in int idAttestationFlags, out KeymasterCertificateChain attestationChain);
     boolean setKeyPairCertificate(in ComponentName who, in String callerPackage, in String alias,
             in byte[] certBuffer, in byte[] certChainBuffer, boolean isUserSelectable);
     void choosePrivateKeyAlias(int uid, in Uri uri, in String alias, IBinder aliasCallback);
diff --git a/keystore/java/android/security/IKeyChainService.aidl b/keystore/java/android/security/IKeyChainService.aidl
index 7c7417d..5a8fa07 100644
--- a/keystore/java/android/security/IKeyChainService.aidl
+++ b/keystore/java/android/security/IKeyChainService.aidl
@@ -34,7 +34,8 @@
     void setUserSelectable(String alias, boolean isUserSelectable);
 
     boolean generateKeyPair(in String algorithm, in ParcelableKeyGenParameterSpec spec);
-    boolean attestKey(in String alias, in byte[] challenge, out KeymasterCertificateChain chain);
+    boolean attestKey(in String alias, in byte[] challenge, in int[] idAttestationFlags,
+            out KeymasterCertificateChain chain);
     boolean setKeyPairCertificate(String alias, in byte[] userCert, in byte[] certChain);
 
     // APIs used by CertInstaller and DevicePolicyManager
diff --git a/keystore/java/android/security/keystore/AttestationUtils.java b/keystore/java/android/security/keystore/AttestationUtils.java
index 0811100..efee8b4 100644
--- a/keystore/java/android/security/keystore/AttestationUtils.java
+++ b/keystore/java/android/security/keystore/AttestationUtils.java
@@ -99,48 +99,35 @@
         }
     }
 
-    /**
-     * Performs attestation of the device's identifiers. This method returns a certificate chain
-     * whose first element contains the requested device identifiers in an extension. The device's
-     * manufacturer, model, brand, device and product are always also included in the attestation.
-     * If the device supports attestation in secure hardware, the chain will be rooted at a
-     * trustworthy CA key. Otherwise, the chain will be rooted at an untrusted certificate. See
-     * <a href="https://developer.android.com/training/articles/security-key-attestation.html">
-     * Key Attestation</a> for the format of the certificate extension.
-     * <p>
-     * Attestation will only be successful when all of the following are true:
-     * 1) The device has been set up to support device identifier attestation at the factory.
-     * 2) The user has not permanently disabled device identifier attestation.
-     * 3) You have permission to access the device identifiers you are requesting attestation for.
-     * <p>
-     * For privacy reasons, you cannot distinguish between (1) and (2). If attestation is
-     * unsuccessful, the device may not support it in general or the user may have permanently
-     * disabled it.
-     *
-     * @param context the context to use for retrieving device identifiers.
-     * @param idTypes the types of device identifiers to attest.
-     * @param attestationChallenge a blob to include in the certificate alongside the device
-     * identifiers.
-     *
-     * @return a certificate chain containing the requested device identifiers in the first element
-     *
-     * @exception SecurityException if you are not permitted to obtain an attestation of the
-     * device's identifiers.
-     * @exception DeviceIdAttestationException if the attestation operation fails.
-     */
-    @RequiresPermission(Manifest.permission.READ_PRIVILEGED_PHONE_STATE)
-    @NonNull public static X509Certificate[] attestDeviceIds(Context context,
-            @NonNull int[] idTypes, @NonNull byte[] attestationChallenge) throws
+    @NonNull private static KeymasterArguments prepareAttestationArgumentsForDeviceId(
+            Context context, @NonNull int[] idTypes, @NonNull byte[] attestationChallenge) throws
             DeviceIdAttestationException {
-        // Check method arguments, retrieve requested device IDs and prepare attestation arguments.
+        // Verify that device ID attestation types are provided.
         if (idTypes == null) {
             throw new NullPointerException("Missing id types");
         }
+
+        return prepareAttestationArguments(context, idTypes, attestationChallenge);
+    }
+
+    /**
+     * Prepares Keymaster Arguments with attestation data.
+     * @hide should only be used by KeyChain.
+     */
+    @NonNull public static KeymasterArguments prepareAttestationArguments(Context context,
+            @NonNull int[] idTypes, @NonNull byte[] attestationChallenge) throws
+            DeviceIdAttestationException {
+        // Check method arguments, retrieve requested device IDs and prepare attestation arguments.
         if (attestationChallenge == null) {
             throw new NullPointerException("Missing attestation challenge");
         }
         final KeymasterArguments attestArgs = new KeymasterArguments();
         attestArgs.addBytes(KeymasterDefs.KM_TAG_ATTESTATION_CHALLENGE, attestationChallenge);
+        // Return early if the caller did not request any device identifiers to be included in the
+        // attestation record.
+        if (idTypes == null) {
+            return attestArgs;
+        }
         final Set<Integer> idTypesSet = new ArraySet<>(idTypes.length);
         for (int idType : idTypes) {
             idTypesSet.add(idType);
@@ -191,6 +178,44 @@
                 Build.MANUFACTURER.getBytes(StandardCharsets.UTF_8));
         attestArgs.addBytes(KeymasterDefs.KM_TAG_ATTESTATION_ID_MODEL,
                 Build.MODEL.getBytes(StandardCharsets.UTF_8));
+        return attestArgs;
+    }
+
+    /**
+     * Performs attestation of the device's identifiers. This method returns a certificate chain
+     * whose first element contains the requested device identifiers in an extension. The device's
+     * manufacturer, model, brand, device and product are always also included in the attestation.
+     * If the device supports attestation in secure hardware, the chain will be rooted at a
+     * trustworthy CA key. Otherwise, the chain will be rooted at an untrusted certificate. See
+     * <a href="https://developer.android.com/training/articles/security-key-attestation.html">
+     * Key Attestation</a> for the format of the certificate extension.
+     * <p>
+     * Attestation will only be successful when all of the following are true:
+     * 1) The device has been set up to support device identifier attestation at the factory.
+     * 2) The user has not permanently disabled device identifier attestation.
+     * 3) You have permission to access the device identifiers you are requesting attestation for.
+     * <p>
+     * For privacy reasons, you cannot distinguish between (1) and (2). If attestation is
+     * unsuccessful, the device may not support it in general or the user may have permanently
+     * disabled it.
+     *
+     * @param context the context to use for retrieving device identifiers.
+     * @param idTypes the types of device identifiers to attest.
+     * @param attestationChallenge a blob to include in the certificate alongside the device
+     * identifiers.
+     *
+     * @return a certificate chain containing the requested device identifiers in the first element
+     *
+     * @exception SecurityException if you are not permitted to obtain an attestation of the
+     * device's identifiers.
+     * @exception DeviceIdAttestationException if the attestation operation fails.
+     */
+    @RequiresPermission(Manifest.permission.READ_PRIVILEGED_PHONE_STATE)
+    @NonNull public static X509Certificate[] attestDeviceIds(Context context,
+            @NonNull int[] idTypes, @NonNull byte[] attestationChallenge) throws
+            DeviceIdAttestationException {
+        final KeymasterArguments attestArgs = prepareAttestationArgumentsForDeviceId(
+                context, idTypes, attestationChallenge);
 
         // Perform attestation.
         final KeymasterCertificateChain outChain = new KeymasterCertificateChain();
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/BaseIDevicePolicyManager.java b/services/devicepolicy/java/com/android/server/devicepolicy/BaseIDevicePolicyManager.java
index e55d4ea..c1e95eb 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/BaseIDevicePolicyManager.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/BaseIDevicePolicyManager.java
@@ -67,7 +67,8 @@
     public void transferOwner(ComponentName admin, ComponentName target, PersistableBundle bundle) {}
 
     public boolean generateKeyPair(ComponentName who, String callerPackage, String algorithm,
-            ParcelableKeyGenParameterSpec keySpec, KeymasterCertificateChain attestationChain) {
+            ParcelableKeyGenParameterSpec keySpec, int idAttestationFlags,
+            KeymasterCertificateChain attestationChain) {
         return false;
     }
 
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
index e5351b4..762aeaf 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
@@ -45,6 +45,10 @@
 import static android.app.admin.DevicePolicyManager.DELEGATION_KEEP_UNINSTALLED_PACKAGES;
 import static android.app.admin.DevicePolicyManager.DELEGATION_PACKAGE_ACCESS;
 import static android.app.admin.DevicePolicyManager.DELEGATION_PERMISSION_GRANT;
+import static android.app.admin.DevicePolicyManager.ID_TYPE_BASE_INFO;
+import static android.app.admin.DevicePolicyManager.ID_TYPE_IMEI;
+import static android.app.admin.DevicePolicyManager.ID_TYPE_MEID;
+import static android.app.admin.DevicePolicyManager.ID_TYPE_SERIAL;
 import static android.app.admin.DevicePolicyManager.LEAVE_ALL_SYSTEM_APPS_ENABLED;
 import static android.app.admin.DevicePolicyManager.PASSWORD_QUALITY_COMPLEX;
 import static android.app.admin.DevicePolicyManager.PASSWORD_QUALITY_UNSPECIFIED;
@@ -217,8 +221,10 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.Date;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map.Entry;
+import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -5052,17 +5058,82 @@
         return false;
     }
 
+    private void enforceIsDeviceOwnerOrCertInstallerOfDeviceOwner(
+            ComponentName who, String callerPackage, int callerUid) throws SecurityException {
+        if (who == null) {
+            if (!mOwners.hasDeviceOwner()) {
+                throw new SecurityException("Not in Device Owner mode.");
+            }
+            if (UserHandle.getUserId(callerUid) != mOwners.getDeviceOwnerUserId()) {
+                throw new SecurityException("Caller not from device owner user");
+            }
+            if (!isCallerDelegate(callerPackage, DELEGATION_CERT_INSTALL)) {
+                throw new SecurityException("Caller with uid " + mInjector.binderGetCallingUid() +
+                        "has no permission to generate keys.");
+            }
+        } else {
+            // Caller provided - check it is the device owner.
+            synchronized (this) {
+                getActiveAdminForCallerLocked(who, DeviceAdminInfo.USES_POLICY_DEVICE_OWNER);
+            }
+        }
+    }
+
+    @VisibleForTesting
+    public static int[] translateIdAttestationFlags(
+            int idAttestationFlags) {
+        Map<Integer, Integer> idTypeToAttestationFlag = new HashMap();
+        idTypeToAttestationFlag.put(ID_TYPE_SERIAL, AttestationUtils.ID_TYPE_SERIAL);
+        idTypeToAttestationFlag.put(ID_TYPE_IMEI, AttestationUtils.ID_TYPE_IMEI);
+        idTypeToAttestationFlag.put(ID_TYPE_MEID, AttestationUtils.ID_TYPE_MEID);
+
+        int numFlagsSet = Integer.bitCount(idAttestationFlags);
+        // No flags are set - return null to indicate no device ID attestation information should
+        // be included in the attestation record.
+        if (numFlagsSet == 0) {
+            return null;
+        }
+
+        // If the ID_TYPE_BASE_INFO is set, make sure that a non-null array is returned, even if
+        // no other flag is set. That will lead to inclusion of general device make data in the
+        // attestation record, but no specific device identifiers.
+        if ((idAttestationFlags & ID_TYPE_BASE_INFO) != 0) {
+            numFlagsSet -= 1;
+            idAttestationFlags = idAttestationFlags & (~ID_TYPE_BASE_INFO);
+        }
+
+        int[] attestationUtilsFlags = new int[numFlagsSet];
+        int i = 0;
+        for (Integer idType: idTypeToAttestationFlag.keySet()) {
+            if ((idType & idAttestationFlags) != 0) {
+                attestationUtilsFlags[i++] = idTypeToAttestationFlag.get(idType);
+            }
+        }
+
+        return attestationUtilsFlags;
+    }
+
     @Override
     public boolean generateKeyPair(ComponentName who, String callerPackage, String algorithm,
             ParcelableKeyGenParameterSpec parcelableKeySpec,
+            int idAttestationFlags,
             KeymasterCertificateChain attestationChain) {
-        enforceCanManageScope(who, callerPackage, DeviceAdminInfo.USES_POLICY_PROFILE_OWNER,
-                DELEGATION_CERT_INSTALL);
+        // Get attestation flags, if any.
+        final int[] attestationUtilsFlags = translateIdAttestationFlags(idAttestationFlags);
+        final boolean deviceIdAttestationRequired = attestationUtilsFlags != null;
+        final int callingUid = mInjector.binderGetCallingUid();
+
+        if (deviceIdAttestationRequired && attestationUtilsFlags.length > 0) {
+            enforceIsDeviceOwnerOrCertInstallerOfDeviceOwner(who, callerPackage, callingUid);
+        } else {
+            enforceCanManageScope(who, callerPackage, DeviceAdminInfo.USES_POLICY_PROFILE_OWNER,
+                    DELEGATION_CERT_INSTALL);
+        }
         final KeyGenParameterSpec keySpec = parcelableKeySpec.getSpec();
-        if (TextUtils.isEmpty(keySpec.getKeystoreAlias())) {
+        final String alias = keySpec.getKeystoreAlias();
+        if (TextUtils.isEmpty(alias)) {
             throw new IllegalArgumentException("Empty alias provided.");
         }
-        final String alias = keySpec.getKeystoreAlias();
         // As the caller will be granted access to the key, ensure no UID was specified, as
         // it will not have the desired effect.
         if (keySpec.getUid() != KeyStore.UID_SELF) {
@@ -5070,7 +5141,10 @@
             return false;
         }
 
-        final int callingUid = mInjector.binderGetCallingUid();
+        if (deviceIdAttestationRequired && (keySpec.getAttestationChallenge() == null)) {
+            throw new IllegalArgumentException(
+                    "Requested Device ID attestation but challenge is empty.");
+        }
 
         final UserHandle userHandle = mInjector.binderGetCallingUserHandle();
         final long id = mInjector.binderClearCallingIdentity();
@@ -5103,7 +5177,7 @@
                 final byte[] attestationChallenge = keySpec.getAttestationChallenge();
                 if (attestationChallenge != null) {
                     final boolean attestationResult = keyChain.attestKey(
-                            alias, attestationChallenge, attestationChain);
+                            alias, attestationChallenge, attestationUtilsFlags, attestationChain);
                     if (!attestationResult) {
                         Log.e(LOG_TAG, String.format(
                                 "Attestation for %s failed, deleting key.", alias));
diff --git a/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java b/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java
index 60783db..58ac7d2 100644
--- a/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java
+++ b/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java
@@ -17,6 +17,10 @@
 
 import static android.app.admin.DevicePolicyManager.DELEGATION_APP_RESTRICTIONS;
 import static android.app.admin.DevicePolicyManager.DELEGATION_CERT_INSTALL;
+import static android.app.admin.DevicePolicyManager.ID_TYPE_BASE_INFO;
+import static android.app.admin.DevicePolicyManager.ID_TYPE_IMEI;
+import static android.app.admin.DevicePolicyManager.ID_TYPE_MEID;
+import static android.app.admin.DevicePolicyManager.ID_TYPE_SERIAL;
 import static android.app.admin.DevicePolicyManager.WIPE_EUICC;
 import static android.os.UserManagerInternal.CAMERA_DISABLED_GLOBALLY;
 import static android.os.UserManagerInternal.CAMERA_DISABLED_LOCALLY;
@@ -71,6 +75,7 @@
 import android.platform.test.annotations.Presubmit;
 import android.provider.Settings;
 import android.security.KeyChain;
+import android.security.keystore.AttestationUtils;
 import android.telephony.TelephonyManager;
 import android.test.MoreAsserts;
 import android.test.suitebuilder.annotation.SmallTest;
@@ -4459,6 +4464,47 @@
         });
     }
 
+    private void assertAttestationFlags(int attestationFlags, int[] expectedFlags) {
+        int[] gotFlags = DevicePolicyManagerService.translateIdAttestationFlags(attestationFlags);
+        Arrays.sort(gotFlags);
+        Arrays.sort(expectedFlags);
+        assertTrue(Arrays.equals(expectedFlags, gotFlags));
+    }
+
+    public void testTranslationOfIdAttestationFlag() {
+        int[] allIdTypes = new int[]{ID_TYPE_SERIAL, ID_TYPE_IMEI, ID_TYPE_MEID};
+        int[] correspondingAttUtilsTypes = new int[]{
+            AttestationUtils.ID_TYPE_SERIAL, AttestationUtils.ID_TYPE_IMEI,
+            AttestationUtils.ID_TYPE_MEID};
+
+        // Test translation of zero flags
+        assertNull(DevicePolicyManagerService.translateIdAttestationFlags(0));
+
+        // Test translation of the ID_TYPE_BASE_INFO flag, which should yield an empty, but
+        // non-null array
+        assertAttestationFlags(ID_TYPE_BASE_INFO, new int[] {});
+
+        // Test translation of a single flag
+        assertAttestationFlags(ID_TYPE_BASE_INFO | ID_TYPE_SERIAL,
+                new int[] {AttestationUtils.ID_TYPE_SERIAL});
+        assertAttestationFlags(ID_TYPE_SERIAL, new int[] {AttestationUtils.ID_TYPE_SERIAL});
+
+        // Test translation of two flags
+        assertAttestationFlags(ID_TYPE_SERIAL | ID_TYPE_IMEI,
+                new int[] {AttestationUtils.ID_TYPE_IMEI, AttestationUtils.ID_TYPE_SERIAL});
+        assertAttestationFlags(ID_TYPE_BASE_INFO | ID_TYPE_MEID | ID_TYPE_SERIAL,
+                new int[] {AttestationUtils.ID_TYPE_MEID, AttestationUtils.ID_TYPE_SERIAL});
+
+        // Test translation of all three flags
+        assertAttestationFlags(ID_TYPE_SERIAL | ID_TYPE_IMEI | ID_TYPE_MEID,
+                new int[] {AttestationUtils.ID_TYPE_IMEI, AttestationUtils.ID_TYPE_SERIAL,
+                    AttestationUtils.ID_TYPE_MEID});
+        // Test translation of all three flags
+        assertAttestationFlags(ID_TYPE_SERIAL | ID_TYPE_IMEI | ID_TYPE_MEID | ID_TYPE_BASE_INFO,
+                new int[] {AttestationUtils.ID_TYPE_IMEI, AttestationUtils.ID_TYPE_SERIAL,
+                    AttestationUtils.ID_TYPE_MEID});
+    }
+
     private void setUserSetupCompleteForUser(boolean isUserSetupComplete, int userhandle) {
         when(getServices().settings.settingsSecureGetIntForUser(Settings.Secure.USER_SETUP_COMPLETE, 0,
                 userhandle)).thenReturn(isUserSetupComplete ? 1 : 0);