Add support for writing host keys to guest attributes.

2 files changed, 82 insertions, 17 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 80785e9..2fb7eaf 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
@@ -28,9 +28,20 @@ from google_compute_engine import file_utils
 from google_compute_engine import logger
 from google_compute_engine import metadata_watcher
 from google_compute_engine.boto import boto_config
+from google_compute_engine.compat import urlerror
+from google_compute_engine.compat import urlrequest
 from google_compute_engine.instance_setup import instance_config
 
 
+class PutRequest(urlrequest.Request):
+    def get_method(self):
+        return 'PUT'
+
+
+GUEST_ATTRIBUTES_URL = ('http://metadata.google.internal/computeMetadata/v1beta1/'
+                        'instance/guest-attributes/hostkeys/')
+
+
 class InstanceSetup(object):
   """Initialize the instance the first time it boots."""
 
@@ -122,6 +133,8 @@ class InstanceSetup(object):
     Args:
       key_type: string, the type of the SSH key.
       key_dest: string, a file location to store the SSH key.
+    Returns:
+      tuple, key_type and public key string.
     """
     # Create a temporary file to save the created RSA keys.
     with tempfile.NamedTemporaryFile(prefix=key_type, delete=True) as temp:
@@ -140,6 +153,27 @@ class InstanceSetup(object):
 
     file_utils.SetPermissions(key_dest, mode=0o600)
     file_utils.SetPermissions('%s.pub' % key_dest, mode=0o644)
+    with open('%s.pub' % key_dest, 'r') as pk:
+        key_data = pk.read()
+
+    key_values = key_data.split()
+    if len(key_values) < 2:
+        return
+    else:
+        return key_values[0], key_values[1]
+
+  def _WriteHostKeyToGuestAttributes(self, key_type, key_value):
+    """Write a host key to guest attributes, ignoring errors."""
+    headers = {'Metadata-Flavor' : 'Google'}
+    url = GUEST_ATTRIBUTES_URL + key_type
+    req = PutRequest(url, key_value, headers)
+    try:
+      response = urlrequest.urlopen(req)
+      self.logger.debug(response)
+      self.logger.info('Wrote %s host key to guest attributes.' % key_type)
+    except urlerror.HTTPError:
+      self.logger.info('Unable to write %s host key to guest attributes.'
+                       % key_type)
 
   def _StartSshd(self):
     """Initialize the SSH daemon."""
@@ -180,7 +214,9 @@ class InstanceSetup(object):
       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)
+        key_data = self._GenerateSshKey(key_type, key_dest)
+        if key_data:
+            self._WriteHostKeyToGuestAttributes(key_data[0], key_data[1])
       self._StartSshd()
       self.instance_config.SetOption(section, 'instance_id', str(instance_id))
 
diff --git a/packages/python-google-compute-engine/google_compute_engine/instance_setup/tests/instance_setup_test.py b/packages/python-google-compute-engine/google_compute_engine/instance_setup/tests/instance_setup_test.py
index 97a4674..1a5c3fd 100644
--- a/packages/python-google-compute-engine/google_compute_engine/instance_setup/tests/instance_setup_test.py
+++ b/packages/python-google-compute-engine/google_compute_engine/instance_setup/tests/instance_setup_test.py
@@ -18,6 +18,7 @@
 import subprocess
 
 from google_compute_engine.instance_setup import instance_setup
+from google_compute_engine.test_compat import builtin
 from google_compute_engine.test_compat import mock
 from google_compute_engine.test_compat import unittest
 
@@ -231,22 +232,25 @@ class InstanceSetupTest(unittest.TestCase):
     temp_dest = '/tmp/dest'
     mock_tempfile.return_value = mock_tempfile
     mock_tempfile.__enter__.return_value.name = temp_dest
-
-    instance_setup.InstanceSetup._GenerateSshKey(
-        self.mock_setup, key_type, key_dest)
-    expected_calls = [
-        mock.call.tempfile(prefix=key_type, delete=True),
-        mock.call.tempfile.__enter__(),
-        mock.call.tempfile.__exit__(None, None, None),
-        mock.call.logger.info(mock.ANY, key_dest),
-        mock.call.call(
-            ['ssh-keygen', '-t', key_type, '-f', temp_dest, '-N', '', '-q']),
-        mock.call.move(temp_dest, key_dest),
-        mock.call.move('%s.pub' % temp_dest, '%s.pub' % key_dest),
-        mock.call.permissions(key_dest, mode=0o600),
-        mock.call.permissions('%s.pub' % key_dest, mode=0o644),
-    ]
-    self.assertEqual(mocks.mock_calls, expected_calls)
+    mock_open = mock.mock_open()
+
+    with mock.patch('%s.open' % builtin, mock_open, create=False):
+      mock_open().read.side_effect = 'ssh-rsa asdfasdf'
+      instance_setup.InstanceSetup._GenerateSshKey(
+          self.mock_setup, key_type, key_dest)
+      expected_calls = [
+          mock.call.tempfile(prefix=key_type, delete=True),
+          mock.call.tempfile.__enter__(),
+          mock.call.tempfile.__exit__(None, None, None),
+          mock.call.logger.info(mock.ANY, key_dest),
+          mock.call.call(
+              ['ssh-keygen', '-t', key_type, '-f', temp_dest, '-N', '', '-q']),
+          mock.call.move(temp_dest, key_dest),
+          mock.call.move('%s.pub' % temp_dest, '%s.pub' % key_dest),
+          mock.call.permissions(key_dest, mode=0o600),
+          mock.call.permissions('%s.pub' % key_dest, mode=0o644),
+      ]
+      self.assertEqual(mocks.mock_calls, expected_calls)
 
   @mock.patch('google_compute_engine.instance_setup.instance_setup.subprocess.check_call')
   def testGenerateSshKeyProcessError(self, mock_call):
@@ -318,6 +322,29 @@ class InstanceSetupTest(unittest.TestCase):
     instance_setup.InstanceSetup._SetSshHostKeys(self.mock_setup)
     self.mock_instance_config.SetOption.assert_not_called()
 
+
+  @mock.patch('google_compute_engine.instance_setup.instance_setup.PutRequest')
+  @mock.patch('google_compute_engine.instance_setup.instance_setup.urlrequest.urlopen')
+  def testWriteHostKeyToGuestAttributes(self, mock_urlopen, mock_put_request):
+    instance_setup.InstanceSetup._WriteHostKeyToGuestAttributes(
+        self.mock_setup, 'ssh-rsa', 'asdfasdf')
+    expected_calls = [
+        mock.call('ssh-rsa', 'asdfasdf')
+    ]
+    self.mock_logger.info.assert_called_with(
+        'Wrote ssh-rsa host key to guest attributes.')
+
+    mock_urlopen.side_effect = instance_setup.urlerror.HTTPError(
+        'http://foo', 403, 'Forbidden', {}, None)
+    instance_setup.InstanceSetup._WriteHostKeyToGuestAttributes(
+        self.mock_setup, 'ssh-rsa', 'asdfasdf')
+    expected_calls = [
+        mock.call('ssh-rsa', 'asdfasdf')
+    ]
+    self.mock_logger.info.assert_called_with(
+        'Unable to write ssh-rsa host key to guest attributes.')
+
+
   @mock.patch('google_compute_engine.instance_setup.instance_setup.os.listdir')
   def testSetSshHostKeysFirstBoot(self, mock_listdir):
     self.mock_instance_config.GetOptionString.return_value = None
@@ -325,6 +352,7 @@ class InstanceSetupTest(unittest.TestCase):
     mock_instance_id.return_value = '123'
     self.mock_setup._GetInstanceId = mock_instance_id
     mock_generate_key = mock.Mock()
+    mock_generate_key.return_value = ('ssh-rsa', 'asdfasdf')
     self.mock_setup._GenerateSshKey = mock_generate_key
     mock_listdir.return_value = [
         'ssh_config',
@@ -344,6 +372,7 @@ class InstanceSetupTest(unittest.TestCase):
         mock.call('ed25519', '/etc/ssh/ssh_host_ed25519_key'),
         mock.call('rsa', '/etc/ssh/ssh_host_rsa_key'),
     ]
+
     self.assertEqual(sorted(mock_generate_key.mock_calls), expected_calls)
     self.mock_instance_config.SetOption.assert_called_once_with(
         'Instance', 'instance_id', '123')