Generate SSH host keys when none are present. (#510)

Provide an instance config option for specifying which host key types to
generate.

3 files changed, 21 insertions, 8 deletions

diff --git a/google_compute_engine/instance_setup/instance_config.py b/google_compute_engine/instance_setup/instance_config.py
index 56c1109..9827137 100644
--- a/google_compute_engine/instance_setup/instance_config.py
+++ b/google_compute_engine/instance_setup/instance_config.py
@@ -72,6 +72,7 @@ class InstanceConfig(config_manager.ConfigManager):
           'instance_id': '0',
       },
       'InstanceSetup': {
+          'host_key_types': 'ecdsa,ed25519,rsa',
           'optimize_local_ssd': 'true',
           'network_enabled': 'true',
           'set_boto_config': 'true',
diff --git a/google_compute_engine/instance_setup/instance_setup.py b/google_compute_engine/instance_setup/instance_setup.py
index 2c6e3d4..e4791c4 100755
--- a/google_compute_engine/instance_setup/instance_setup.py
+++ b/google_compute_engine/instance_setup/instance_setup.py
@@ -55,7 +55,9 @@ class InstanceSetup(object):
       self.instance_config = instance_config.InstanceConfig(
           logger=self.logger, instance_config_metadata=instance_config_metadata)
       if self.instance_config.GetOptionBool('InstanceSetup', 'set_host_keys'):
-        self._SetSshHostKeys()
+        host_key_types = self.instance_config.GetOptionString(
+            'InstanceSetup', 'host_key_types')
+        self._SetSshHostKeys(host_key_types=host_key_types)
       if self.instance_config.GetOptionBool('InstanceSetup', 'set_boto_config'):
         self._SetupBotoConfig()
     if self.instance_config.GetOptionBool(
@@ -155,13 +157,16 @@ class InstanceSetup(object):
       subprocess.call(['service', 'sshd', 'start'])
       subprocess.call(['service', 'sshd', 'reload'])
 
-  def _SetSshHostKeys(self):
+  def _SetSshHostKeys(self, host_key_types=None):
     """Regenerates SSH host keys when the VM is restarted with a new IP address.
 
     Booting a VM from an image with a known SSH key allows a number of attacks.
     This function will regenerating the host key whenever the IP address
     changes. This applies the first time the instance is booted, and each time
     the disk is used to boot a new instance.
+
+    Args:
+      host_key_types: string, a comma separated list of host key types.
     """
     section = 'Instance'
     instance_id = self._GetInstanceId()
@@ -171,7 +176,9 @@ class InstanceSetup(object):
       file_regex = re.compile(r'ssh_host_(?P<type>[a-z0-9]*)_key\Z')
       key_dir = '/etc/ssh'
       key_files = [f for f in os.listdir(key_dir) if file_regex.match(f)]
-      for key_file in key_files:
+      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):
         key_type = file_regex.match(key_file).group('type')
         key_dest = os.path.join(key_dir, key_file)
         self._GenerateSshKey(key_type, key_dest)
diff --git a/google_compute_engine/instance_setup/tests/instance_setup_test.py b/google_compute_engine/instance_setup/tests/instance_setup_test.py
index 1de8134..bf7c18b 100644
--- a/google_compute_engine/instance_setup/tests/instance_setup_test.py
+++ b/google_compute_engine/instance_setup/tests/instance_setup_test.py
@@ -49,6 +49,7 @@ class InstanceSetupTest(unittest.TestCase):
     mock_watcher.MetadataWatcher.return_value = mock_watcher_instance
     mock_config_instance = mock.Mock()
     mock_config_instance.GetOptionBool.return_value = True
+    mock_config_instance.GetOptionString.return_value = 'type'
     mock_config.InstanceConfig.return_value = mock_config_instance
     mock_setup._GetInstanceConfig.return_value = 'config'
 
@@ -70,7 +71,9 @@ class InstanceSetupTest(unittest.TestCase):
         # Setup for SSH host keys if necessary.
         mock.call.config.InstanceConfig().GetOptionBool(
             'InstanceSetup', 'set_host_keys'),
-        mock.call.setup._SetSshHostKeys(),
+        mock.call.config.InstanceConfig().GetOptionString(
+            'InstanceSetup', 'host_key_types'),
+        mock.call.setup._SetSshHostKeys(host_key_types='type'),
         # Setup for the boto config if necessary.
         mock.call.config.InstanceConfig().GetOptionBool(
             'InstanceSetup', 'set_boto_config'),
@@ -325,7 +328,7 @@ class InstanceSetupTest(unittest.TestCase):
     self.mock_setup._GenerateSshKey = mock_generate_key
     mock_listdir.return_value = [
         'ssh_config',
-        'ssh_host_rsa_key',
+        'ssh_host_dsa_key',
         'ssh_host_dsa_key.pub',
         'ssh_host_ed25519_key',
         'ssh_host_ed25519_key.pub',
@@ -333,13 +336,15 @@ class InstanceSetupTest(unittest.TestCase):
         'ssh_host_rsa_key.pub',
     ]
 
-    instance_setup.InstanceSetup._SetSshHostKeys(self.mock_setup)
+    instance_setup.InstanceSetup._SetSshHostKeys(
+        self.mock_setup, host_key_types='rsa,dsa,abc')
     expected_calls = [
-        mock.call('rsa', '/etc/ssh/ssh_host_rsa_key'),
+        mock.call('abc', '/etc/ssh/ssh_host_abc_key'),
+        mock.call('dsa', '/etc/ssh/ssh_host_dsa_key'),
         mock.call('ed25519', '/etc/ssh/ssh_host_ed25519_key'),
         mock.call('rsa', '/etc/ssh/ssh_host_rsa_key'),
     ]
-    self.assertEqual(mock_generate_key.mock_calls, expected_calls)
+    self.assertEqual(sorted(mock_generate_key.mock_calls), expected_calls)
     self.mock_instance_config.SetOption.assert_called_once_with(
         'Instance', 'instance_id', '123')