diff options
2 files changed, 96 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..e72534c 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,21 @@ 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') +HOSTKEY_NAMESPACE = 'hostkeys' + + class InstanceSetup(object): """Initialize the instance the first time it boots.""" @@ -122,6 +134,9 @@ 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 +155,28 @@ 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: + self.logger.warning('Could not read host key from %s.pub.', key_dest) + 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 = '%s/%s/%s' % (GUEST_ATTRIBUTES_URL, HOSTKEY_NAMESPACE, 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 +217,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..a59ae0c 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,33 @@ 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() + key_file_contents = 'ssh-rsa asdfasdf' + expected_key_data = ('ssh-rsa', 'asdfasdf') + + with mock.patch('%s.open' % builtin, mock_open, create=False): + mock_open().read.return_value = key_file_contents + key_data = 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) + self.assertEqual(key_data, expected_key_data) + + mock_open().read.return_value = '' + key_data = instance_setup.InstanceSetup._GenerateSshKey( + self.mock_setup, key_type, key_dest) + self.assertEqual(key_data, None) @mock.patch('google_compute_engine.instance_setup.instance_setup.subprocess.check_call') def testGenerateSshKeyProcessError(self, mock_call): @@ -318,6 +330,32 @@ 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.urlrequest.urlopen') + @mock.patch('google_compute_engine.instance_setup.instance_setup.PutRequest') + def testWriteHostKeyToGuestAttributes(self, mock_put, mock_urlopen): + key_type = 'ssh-rsa' + key_value = 'asdfasdf' + expected_url = ('http://metadata.google.internal/computeMetadata/v1beta1/' + 'instance/guest-attributes/hostkeys/%s' % key_type) + headers = {'Metadata-Flavor': 'Google'} + + instance_setup.InstanceSetup._WriteHostKeyToGuestAttributes( + self.mock_setup, key_type, key_value) + self.mock_logger.info.assert_called_with( + 'Wrote %s host key to guest attributes.', key_type) + mock_put.assert_called_with(expected_url, key_value, headers) + + mock_urlopen.side_effect = instance_setup.urlerror.HTTPError( + 'http://foo', 403, 'Forbidden', {}, None) + instance_setup.InstanceSetup._WriteHostKeyToGuestAttributes( + self.mock_setup, key_type, key_value) + self.mock_logger.info.assert_called_with( + 'Unable to write %s host key to guest attributes.', key_type) + + def testPutRequest(self): + put_request = instance_setup.PutRequest('http://example.com/') + self.assertEqual(put_request.get_method(), 'PUT') + @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 +363,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 +383,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') |