diff options
author | Rick Wright <rickw@google.com> | 2019-04-29 14:20:07 -0700 |
---|---|---|
committer | Rick Wright <rickw@google.com> | 2019-04-29 14:20:07 -0700 |
commit | 984c8b5ae2092286093e48c8d72462a9990d97cc (patch) | |
tree | a3eb3200077684728593a0ad9cc7066483fd499a /packages/python-google-compute-engine/google_compute_engine | |
parent | fd564de77d393ac6ed078075e888737f0aa57e65 (diff) | |
download | google-compute-image-packages-984c8b5ae2092286093e48c8d72462a9990d97cc.tar.gz |
Add support for writing host keys to guest attributes.
Diffstat (limited to 'packages/python-google-compute-engine/google_compute_engine')
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') |