Use private key to create the ssh public key.

if private key exist but no public key, use the private key to
create the public key.

Bug: b/110856450
Test: ./run_tests.sh && m acloud && rm ~/.ssh/acloud_rsa.pub && acloud create
Change-Id: Ie2b482610b7c1a6331745a8e8e9a57b0363d9092
diff --git a/internal/lib/utils.py b/internal/lib/utils.py
index 35729a6..2cd3ff2 100755
--- a/internal/lib/utils.py
+++ b/internal/lib/utils.py
@@ -35,6 +35,7 @@
 logger = logging.getLogger(__name__)
 
 SSH_KEYGEN_CMD = ["ssh-keygen", "-t", "rsa", "-b", "4096"]
+SSH_KEYGEN_PUB_CMD = ["ssh-keygen", "-y"]
 DEFAULT_RETRY_BACKOFF_FACTOR = 1
 DEFAULT_SLEEP_MULTIPLIER = 0
 
@@ -123,13 +124,8 @@
     return _Wrapper
 
 
-def Retry(retry_checker,
-          max_retries,
-          functor,
-          sleep_multiplier,
-          retry_backoff_factor,
-          *args,
-          **kwargs):
+def Retry(retry_checker, max_retries, functor, sleep_multiplier,
+          retry_backoff_factor, *args, **kwargs):
     """Conditionally retry a function.
 
     Args:
@@ -258,34 +254,54 @@
 def CreateSshKeyPairIfNotExist(private_key_path, public_key_path):
     """Create the ssh key pair if they don't exist.
 
-    Check if the public and private key pairs exist at
-    the given places. If not, create them.
+    Case1. If the private key doesn't exist, we will create both the public key
+           and the private key.
+    Case2. If the private key exists but public key doesn't, we will create the
+           public key by using the private key.
+    Case3. If the public key exists but the private key doesn't, we will create
+           a new private key and overwrite the public key.
 
     Args:
         private_key_path: Path to the private key file.
                           e.g. ~/.ssh/acloud_rsa
         public_key_path: Path to the public key file.
                          e.g. ~/.ssh/acloud_rsa.pub
+
     Raises:
         error.DriverError: If failed to create the key pair.
     """
     public_key_path = os.path.expanduser(public_key_path)
     private_key_path = os.path.expanduser(private_key_path)
-    create_key = (not os.path.exists(public_key_path)
-                  and not os.path.exists(private_key_path))
-    if not create_key:
+    public_key_exist = os.path.exists(public_key_path)
+    private_key_exist = os.path.exists(private_key_path)
+    if public_key_exist and private_key_exist:
         logger.debug(
-            "The ssh private key (%s) or public key (%s) already exist,"
+            "The ssh private key (%s) and public key (%s) already exist,"
             "will not automatically create the key pairs.", private_key_path,
             public_key_path)
         return
-    cmd = SSH_KEYGEN_CMD + ["-C", getpass.getuser(), "-f", private_key_path]
-    logger.info(
-        "The ssh private key (%s) and public key (%s) do not exist, "
-        "automatically creating key pair, calling: %s", private_key_path,
-        public_key_path, " ".join(cmd))
+    key_folder = os.path.dirname(private_key_path)
+    if not os.path.exists(key_folder):
+        os.makedirs(key_folder)
     try:
-        subprocess.check_call(cmd, stdout=sys.stderr, stderr=sys.stdout)
+        if private_key_exist:
+            cmd = SSH_KEYGEN_PUB_CMD + ["-f", private_key_path]
+            with open(public_key_path, 'w') as outfile:
+                stream_content = subprocess.check_output(cmd)
+                outfile.write(
+                    stream_content.rstrip('\n') + " " + getpass.getuser())
+            logger.info(
+                "The ssh public key (%s) do not exist, "
+                "automatically creating public key, calling: %s",
+                public_key_path, " ".join(cmd))
+        else:
+            cmd = SSH_KEYGEN_CMD + [
+                "-C", getpass.getuser(), "-f", private_key_path
+            ]
+            logger.info(
+                "Creating public key from private key (%s) via cmd: %s",
+                private_key_path, " ".join(cmd))
+            subprocess.check_call(cmd, stdout=sys.stderr, stderr=sys.stdout)
     except subprocess.CalledProcessError as e:
         raise errors.DriverError("Failed to create ssh key pair: %s" % str(e))
     except OSError as e: