path: root/packages/python-google-compute-engine/google_compute_engine/instance_setup/instance_setup.py

#!/usr/bin/python
# Copyright 2016 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Run initialization code the first time the instance boots."""

import logging.handlers
import optparse
import os
import re
import shutil
import subprocess
import tempfile

from google_compute_engine import constants
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 distro_name, 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/v1/'
                        'instance/guest-attributes')
HOSTKEY_NAMESPACE = 'hostkeys'


class InstanceSetup(object):
  """Initialize the instance the first time it boots."""

  def __init__(self, debug=False):
    """Constructor.

    Args:
      debug: bool, True if debug output should write to the console.
    """
    self.debug = debug
    facility = logging.handlers.SysLogHandler.LOG_DAEMON
    self.logger = logger.Logger(
        name='instance-setup', debug=self.debug, facility=facility)
    self.watcher = metadata_watcher.MetadataWatcher(logger=self.logger)
    self.metadata_dict = None
    self.instance_config = instance_config.InstanceConfig(logger=self.logger)

    if self.instance_config.GetOptionBool('InstanceSetup', 'network_enabled'):
      self.metadata_dict = self.watcher.GetMetadata()
      instance_config_metadata = self._GetInstanceConfig()
      self.instance_config = instance_config.InstanceConfig(
          logger=self.logger, instance_config_metadata=instance_config_metadata)

      if self.instance_config.GetOptionBool('InstanceSetup', 'set_host_keys'):
        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()

      # machineType is e.g. u'projects/00000000000000/machineTypes/n1-standard-1'
      machineType = self.metadata_dict['instance']['machineType'].split('/')[-1]
      if machineType.startswith("e2-") and 'bsd' not in distro_name:  # Not yet supported on BSD.
        subprocess.call(["sysctl", "vm.overcommit_memory=1"])

    if self.instance_config.GetOptionBool(
        'InstanceSetup', 'optimize_local_ssd'):
      self._RunScript('google_optimize_local_ssd')

    if self.instance_config.GetOptionBool('InstanceSetup', 'set_multiqueue'):
      self._RunScript('google_set_multiqueue')

    try:
      self.instance_config.WriteConfig()
    except (IOError, OSError) as e:
      self.logger.warning(str(e))

  def _GetInstanceConfig(self):
    """Get the instance configuration specified in metadata.

    Returns:
      string, the instance configuration data.
    """
    try:
      instance_data = self.metadata_dict['instance']['attributes']
    except KeyError:
      instance_data = {}
      self.logger.warning('Instance attributes were not found.')

    try:
      project_data = self.metadata_dict['project']['attributes']
    except KeyError:
      project_data = {}
      self.logger.warning('Project attributes were not found.')

    return (instance_data.get('google-instance-configs')
            or project_data.get('google-instance-configs'))

  def _RunScript(self, script):
    """Run a script and log the streamed script output.

    Args:
      script: string, the file location of an executable script.
    """
    process = subprocess.Popen(
        script, shell=True, stderr=subprocess.STDOUT, stdout=subprocess.PIPE)
    while True:
      for line in iter(process.stdout.readline, b''):
        self.logger.info(line.decode('utf-8').rstrip('\n'))
      if process.poll() is not None:
        break

  def _GetInstanceId(self):
    """Get the instance ID for this VM.

    Returns:
      string, the instance ID for the VM.
    """
    try:
      return str(self.metadata_dict['instance']['id'])
    except KeyError:
      self.logger.warning('Instance ID was not found in metadata.')
      return None

  def _GenerateSshKey(self, key_type, key_dest):
    """Generate a new SSH key.

    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:
      temp_key = temp.name

    command = ['ssh-keygen', '-t', key_type, '-f', temp_key, '-N', '', '-q']
    try:
      self.logger.info('Generating SSH key %s.', key_dest)
      subprocess.check_call(command)
    except subprocess.CalledProcessError:
      self.logger.warning('Could not create SSH key %s.', key_dest)
      return

    shutil.move(temp_key, key_dest)
    shutil.move('%s.pub' % temp_key, '%s.pub' % key_dest)

    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)
    key_value = key_value.encode('utf-8')
    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."""
    # Exit as early as possible.
    # Instance setup systemd scripts block sshd from starting.
    if os.path.exists(constants.LOCALBASE + '/bin/systemctl'):
      return
    elif (os.path.exists('/etc/init.d/ssh')
          or os.path.exists('/etc/init/ssh.conf')):
      subprocess.call(['service', 'ssh', 'start'])
      subprocess.call(['service', 'ssh', 'reload'])
    elif (os.path.exists('/etc/init.d/sshd')
          or os.path.exists('/etc/init/sshd.conf')):
      subprocess.call(['service', 'sshd', 'start'])
      subprocess.call(['service', 'sshd', 'reload'])

  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.
    """
    instance_id = self._GetInstanceId()
    if instance_id != self.instance_config.GetOptionString(
        'instance', '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_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):
        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()
      self.instance_config.SetOption('Instance', 'instance_id', str(instance_id))

  def _GetNumericProjectId(self):
    """Get the numeric project ID.

    Returns:
      string, the numeric project ID.
    """
    try:
      return str(self.metadata_dict['project']['numericProjectId'])
    except KeyError:
      self.logger.warning('Numeric project ID was not found in metadata.')
      return None

  def _SetupBotoConfig(self):
    """Set the boto config so GSUtil works with provisioned service accounts."""
    project_id = self._GetNumericProjectId()
    try:
      boto_config.BotoConfig(project_id, debug=self.debug)
    except (IOError, OSError) as e:
      self.logger.warning(str(e))


def main():
  parser = optparse.OptionParser()
  parser.add_option(
      '-d', '--debug', action='store_true', dest='debug',
      help='print debug output to the console.')
  (options, _) = parser.parse_args()
  InstanceSetup(debug=bool(options.debug))


if __name__ == '__main__':
  main()