From 54248185371a5017913e15451858959a78c9621d Mon Sep 17 00:00:00 2001
From: Dexter Rivera <riverade@google.com>
Date: Wed, 4 Mar 2020 16:00:12 -0800
Subject: [PATCH] publish hostkeys from stateful partition

---
 .../instance_setup/instance_setup.py          | 38 ++++++++++++++-----
 1 file changed, 28 insertions(+), 10 deletions(-)

diff --git a/packages/python-google-compute-engine/google_compute_engine/instance_setup/instance_setup.py b/packages/python-google-compute-engine/google_compute_engine/instance_setup/instance_setup.py
index cb1a2a6..531b7d4 100755
--- a/packages/python-google-compute-engine/google_compute_engine/instance_setup/instance_setup.py
+++ b/packages/python-google-compute-engine/google_compute_engine/instance_setup/instance_setup.py
@@ -207,23 +207,41 @@ class InstanceSetup(object):
     """
     section = 'Instance'
     instance_id = self._GetInstanceId()
-    if instance_id != self.instance_config.GetOptionString(
-        section, 'instance_id'):
+    prev_instance_id = None
+
+    instance_id_file = '/mnt/stateful_partition/.instance_id'
+    if os.path.isfile(instance_id_file):
+      with open(instance_id_file, 'rb') as f:
+        prev_instance_id = f.read().strip()
+
+    if not prev_instance_id or prev_instance_id != instance_id:
       self.logger.info('Generating SSH host keys for instance %s.', instance_id)
       file_regex = re.compile(r'ssh_host_(?P<type>[a-z0-9]*)_key\Z')
-      key_dir = '/etc/ssh'
+      key_dir = '/mnt/stateful_partition/etc/ssh'
       key_files = [f for f in os.listdir(key_dir) if file_regex.match(f)]
-      key_types = host_key_types.split(',') if host_key_types else []
-      key_types_files = ['ssh_host_%s_key' % key_type for key_type in key_types]
-      for key_file in set(key_files) | set(key_types_files):
+      for key_file in set(key_files):
         key_type = file_regex.match(key_file).group('type')
         key_dest = os.path.join(key_dir, key_file)
-        key_data = self._GenerateSshKey(key_type, key_dest)
-        if key_data:
-          self._WriteHostKeyToGuestAttributes(key_data[0], key_data[1])
-      self._StartSshd()
+        if not prev_instance_id:
+          # On first boot, the hostkeys will be generated by a script
+          # that runs before sshd is started, so publish those keys
+          with open('%s.pub' % key_dest, 'rb') as pk:
+            key_data = pk.read().split()
+            if len(key_data) < 2:
+              self.logger.warning(
+                  'Could not read host key from %s.pub. Unable to publish '
+                  '%s key to instance guest attributes', key_dest, key_type)
+              continue
+        else:
+          # If the instance id changes, need to regenerate the hostkeys
+          key_data = self._GenerateSshKey(key_type, key_dest)
+        self._WriteHostKeyToGuestAttributes(key_data[0], key_data[1])
       self.instance_config.SetOption(section, 'instance_id', str(instance_id))
 
+      # Write the instance_id to the stateful partition
+      with open(instance_id_file, 'wb') as f:
+        f.write(instance_id)
+
   def _GetNumericProjectId(self):
     """Get the numeric project ID.
 
-- 
2.25.0.265.gbab2e86ba0-goog