1 files changed, 155 insertions, 0 deletions
diff --git a/tools/tradefed-host/src/com/android/afwtest/tradefed/targetprep/AfwTestEncryptDevice.java b/tools/tradefed-host/src/com/android/afwtest/tradefed/targetprep/AfwTestEncryptDevice.java
new file mode 100644
index 0000000..7de4ace
--- /dev/null
+++ b/tools/tradefed-host/src/com/android/afwtest/tradefed/targetprep/AfwTestEncryptDevice.java
@@ -0,0 +1,155 @@
+/*
+ * Copyright (C) 2016 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 com.android.afwtest.tradefed.targetprep;
+
+import com.android.tradefed.build.IBuildInfo;
+import com.android.tradefed.config.Option;
+import com.android.tradefed.config.OptionClass;
+import com.android.tradefed.device.DeviceNotAvailableException;
+import com.android.tradefed.device.ITestDevice;
+import com.android.tradefed.log.LogUtil.CLog;
+import com.android.tradefed.targetprep.ITargetPreparer;
+import com.android.tradefed.targetprep.TargetSetupError;
+import com.android.tradefed.device.CollectingOutputReceiver;
+import com.android.ddmlib.NullOutputReceiver;
+
+
+import java.util.concurrent.TimeUnit;
+
+/**
+ * A {@link ITargetPreparer} that encrypts the testing device.
+ *
+ * <p>Requires a device where 'adb root' is possible, typically a userdebug build type.</p>
+ *
+ * <p>If the device is already encrypted, nothing will be done. </p>
+ */
+@OptionClass(alias = "afw-test-encrypt-device")
+public class AfwTestEncryptDevice extends AfwTestTargetPreparer implements ITargetPreparer {
+
+    @Option(name = "inplace",
+            description = "Whether to encrypt the device inplace or by wipe")
+    private Boolean mInplace = true;
+
+    @Option(name = "timeout-secs",
+            description = "Timeout in seconds.")
+    private Long mTimeoutSecs = TimeUnit.MINUTES.toSeconds(5);
+
+    @Option(name = "attempts",
+            description = "Number of attempts to encrypt the device")
+    private int mAttempts = 3;
+
+    /** Adb command timeout, in milliseconds. */
+    private static final int CMD_TIMEOUT = (int)TimeUnit.MINUTES.toMillis(2);
+
+    /** The password for encrypting and decrypting the device. */
+    private static final String ENCRYPTION_PASSWORD = "android";
+
+    /** Encrypting with inplace can take up to 2 hours. */
+    private static final int ENCRYPTION_INPLACE_TIMEOUT_MIN = 2 * 60;
+
+    /** Encrypting with wipe can take up to 20 minutes. */
+    private static final long ENCRYPTION_WIPE_TIMEOUT_MIN = 20;
+
+    /**
+     * {@inheritDoc}
+     */
+    @Override
+    public void setUp(ITestDevice device, IBuildInfo buildInfo) throws TargetSetupError,
+            DeviceNotAvailableException {
+
+        for (int i = 0; i < mAttempts; ++i) {
+            CLog.i(String.format("Encrypting device %s: #%d", device.getSerialNumber(), i+1));
+
+            // Exit from loop if the device is encrypted successfully
+            if (encryptDevice(device)) {
+                break;
+            }
+
+            // The device may become unavailable, wait for it.
+            device.waitForDeviceAvailable(TimeUnit.SECONDS.toMillis(10));
+        }
+
+        // Fail the target preparer if encryption failed
+        if (!device.isDeviceEncrypted()) {
+            throw new TargetSetupError(
+                    "Failed to encrypt device " + device.getSerialNumber());
+        }
+
+        device.waitForDeviceAvailable(TimeUnit.SECONDS.toMillis(mTimeoutSecs) * getTimeoutSize());
+        CLog.i(String.format("Device %s is encrypted successfully", device.getSerialNumber()));
+
+        postDeviceEncryption(device);
+    }
+
+    /**
+     * Encrypts the device.
+     *
+     * @param device test device
+     * @return {@code true} if device was encrypted successfully; {@code false} otherwise
+     */
+    protected boolean encryptDevice(ITestDevice device) throws DeviceNotAvailableException {
+        if (!device.isEncryptionSupported()) {
+            throw new UnsupportedOperationException(String.format("Can't encrypt device %s: "
+                    + "encryption not supported", device.getSerialNumber()));
+        }
+
+        if (device.isDeviceEncrypted()) {
+            CLog.d("Device %s is already encrypted, skipping...", device.getSerialNumber());
+            return true;
+        }
+
+        String encryptMethod;
+        long timeout;
+        if (mInplace) {
+            encryptMethod = "inplace";
+            timeout = ENCRYPTION_INPLACE_TIMEOUT_MIN;
+        } else {
+            encryptMethod = "wipe";
+            timeout = ENCRYPTION_WIPE_TIMEOUT_MIN;
+        }
+
+        CLog.i("Encrypting device %s via %s", device.getSerialNumber(), encryptMethod);
+
+        // enable crypto takes one of the following formats:
+        // cryptfs enablecrypto <wipe|inplace> <passwd>
+        // cryptfs enablecrypto <wipe|inplace> default|password|pin|pattern [passwd]
+        // Try the first one first, if it outputs "500 <process_id> Usage: ...", try the second.
+        CollectingOutputReceiver receiver = new CollectingOutputReceiver();
+        String command = String.format("vdc cryptfs enablecrypto %s \"%s\"", encryptMethod,
+                ENCRYPTION_PASSWORD);
+        device.executeShellCommand(command, receiver, timeout, TimeUnit.MINUTES, 1);
+        if (receiver.getOutput().split(":")[0].matches("500 \\d+ Usage")) {
+            command = String.format("vdc cryptfs enablecrypto %s default", encryptMethod);
+            device.executeShellCommand(command,
+                    new NullOutputReceiver(), timeout, TimeUnit.MINUTES, 1);
+        }
+
+        device.waitForDeviceNotAvailable(CMD_TIMEOUT);
+        device.waitForDeviceOnline();
+
+        return device.isDeviceEncrypted();
+    }
+
+    /**
+     * Actions after encryption.
+     *
+     * @param device test device
+     */
+    protected void postDeviceEncryption(ITestDevice device) throws DeviceNotAvailableException {
+        // By default do nothing.
+    }
+}