[Acloud] Automatically creating ssh key pair

Cherry-pick cl/148800637
BUG: 32411892
TEST: Manually tested the google3 version.

Change-Id: I4d0b8649463185c5dee831a7db003dbb150392e0
diff --git a/public/acloud_main.py b/public/acloud_main.py
index 4a7a7b1..d00cd91 100755
--- a/public/acloud_main.py
+++ b/public/acloud_main.py
@@ -33,7 +33,8 @@
      client_secret: <client secret for the client id>
 
      # Optional
-     ssh_private_key_path: ""
+     ssh_private_key_path: "~/.ssh/acloud_rsa"
+     ssh_public_key_path: "~/.ssh/acloud_rsa.pub"
      orientation: "portrait"
      resolution: "800x1280x32x213"
      network: "default"
@@ -80,7 +81,7 @@
 CMD_CREATE = "create"
 CMD_DELETE = "delete"
 CMD_CLEANUP = "cleanup"
-CMD_SSHKEY = "sshkey"
+CMD_SSHKEY = "project_sshkey"
 
 
 def _ParseArgs(args):
@@ -195,7 +196,7 @@
         "images that are older than |expiration_mins|.")
     subparser_list.append(cleanup_parser)
 
-    # Command "sshkey"
+    # Command "project_sshkey"
     sshkey_parser = subparsers.add_parser(CMD_SSHKEY)
     sshkey_parser.required = False
     sshkey_parser.set_defaults(which=CMD_SSHKEY)
@@ -211,7 +212,8 @@
         type=str,
         dest="ssh_rsa_path",
         required=True,
-        help="Absolute path to the file that contains the public rsa key.")
+        help="Absolute path to the file that contains the public rsa key "
+             "that will be added as project-wide ssh key.")
     subparser_list.append(sshkey_parser)
 
     # Add common arguments.
diff --git a/public/device_driver.py b/public/device_driver.py
index 0f2acca..ae7e3c8 100755
--- a/public/device_driver.py
+++ b/public/device_driver.py
@@ -322,6 +322,36 @@
         utils.MakeTarFile(src_dict, output_file)
 
 
+def _CreateSshKeyPairIfNecessary(cfg):
+    """Create ssh key pair if necessary.
+
+    Args:
+        cfg: An Acloudconfig instance.
+
+    Raises:
+        error.DriverError: If it falls into an unexpected condition.
+    """
+    if not cfg.ssh_public_key_path:
+        logger.warning("ssh_public_key_path is not specified in acloud config. "
+                       "Project-wide public key will "
+                       "be used when creating AVD instances. "
+                       "Please ensure you have the correct private half of "
+                       "a project-wide public key if you want to ssh into the "
+                       "instances after creation.")
+    elif cfg.ssh_public_key_path and not cfg.ssh_private_key_path:
+        logger.warning("Only ssh_public_key_path is specified in acloud config,"
+                       " but ssh_private_key_path is missing. "
+                       "Please ensure you have the correct private half "
+                       "if you want to ssh into the instances after creation.")
+    elif cfg.ssh_public_key_path and cfg.ssh_private_key_path:
+        utils.CreateSshKeyPairIfNotExist(
+                cfg.ssh_private_key_path, cfg.ssh_public_key_path)
+    else:
+        # Should never reach here.
+        raise errors.DriverError(
+                "Unexpected error in _CreateSshKeyPairIfNecessary")
+
+
 def CreateAndroidVirtualDevices(cfg,
                                 build_target=None,
                                 build_id=None,
@@ -357,6 +387,7 @@
     compute_client = android_compute_client.AndroidComputeClient(cfg,
                                                                  credentials)
     try:
+        _CreateSshKeyPairIfNecessary(cfg)
         device_pool = AndroidVirtualDevicePool(cfg)
         device_pool.CreateDevices(
             num,
diff --git a/public/device_driver_test.py b/public/device_driver_test.py
index eb0a64f..5155f09 100644
--- a/public/device_driver_test.py
+++ b/public/device_driver_test.py
@@ -64,6 +64,8 @@
             4: "extradisk-image-4gb",
             10: "extradisk-image-10gb"
         }
+        cfg.ssh_private_key_path = ""
+        cfg.ssh_public_key_path = ""
 
         return cfg