juju/provisioner.py

   1 from .client import client
   2
   3 import paramiko
   4 import os
   5 import re
   6 import tempfile
   7 import shlex
   8 from subprocess import (
   9     CalledProcessError,
  10 )
  11 import uuid
  12
  13
  14 arches = [
  15     [re.compile(r"amd64|x86_64"), "amd64"],
  16     [re.compile(r"i?[3-9]86"), "i386"],
  17     [re.compile(r"(arm$)|(armv.*)"), "armhf"],
  18     [re.compile(r"aarch64"), "arm64"],
  19     [re.compile(r"ppc64|ppc64el|ppc64le"), "ppc64el"],
  20     [re.compile(r"s390x?"), "s390x"],
  21
  22 ]
  23
  24
  25 def normalize_arch(rawArch):
  26     """Normalize the architecture string."""
  27     for arch in arches:
  28         if arch[0].match(rawArch):
  29             return arch[1]
  30
  31
  32 DETECTION_SCRIPT = """#!/bin/bash
  33 set -e
  34 os_id=$(grep '^ID=' /etc/os-release | tr -d '"' | cut -d= -f2)
  35 if [ "$os_id" = 'centos' ]; then
  36   os_version=$(grep '^VERSION_ID=' /etc/os-release | tr -d '"' | cut -d= -f2)
  37   echo "centos$os_version"
  38 else
  39   lsb_release -cs
  40 fi
  41 uname -m
  42 grep MemTotal /proc/meminfo
  43 cat /proc/cpuinfo
  44 """
  45
  46 INITIALIZE_UBUNTU_SCRIPT = """set -e
  47 (id ubuntu &> /dev/null) || useradd -m ubuntu -s /bin/bash
  48 umask 0077
  49 temp=$(mktemp)
  50 echo 'ubuntu ALL=(ALL) NOPASSWD:ALL' > $temp
  51 install -m 0440 $temp /etc/sudoers.d/90-juju-ubuntu
  52 rm $temp
  53 su ubuntu -c 'install -D -m 0600 /dev/null ~/.ssh/authorized_keys'
  54 export authorized_keys="{}"
  55 if [ ! -z "$authorized_keys" ]; then
  56     su ubuntu -c 'echo $authorized_keys >> ~/.ssh/authorized_keys'
  57 fi
  58 """
  59
  60
  61 class SSHProvisioner:
  62     """Provision a manually created machine via SSH."""
  63     user = ""
  64     host = ""
  65     private_key_path = ""
  66
  67     def __init__(self, user, host, private_key_path):
  68         self.host = host
  69         self.user = user
  70         self.private_key_path = private_key_path
  71
  72     def _get_ssh_client(self, host, user, key):
  73         """Return a connected Paramiko ssh object.
  74
  75         :param str host: The host to connect to.
  76         :param str user: The user to connect as.
  77         :param str key: The private key to authenticate with.
  78
  79         :return: object: A paramiko.SSHClient
  80         :raises: :class:`paramiko.ssh_exception.SSHException` if the
  81             connection failed
  82         """
  83
  84         ssh = paramiko.SSHClient()
  85         ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
  86
  87         pkey = None
  88
  89         # Read the private key into a paramiko.RSAKey
  90         if os.path.exists(key):
  91             with open(key, 'r') as f:
  92                 pkey = paramiko.RSAKey.from_private_key(f)
  93
  94         #######################################################################
  95         # There is a bug in some versions of OpenSSH 4.3 (CentOS/RHEL5) where #
  96         # the server may not send the SSH_MSG_USERAUTH_BANNER message except  #
  97         # when responding to an auth_none request. For example, paramiko will #
  98         # attempt to use password authentication when a password is set, but  #
  99         # the server could deny that, instead requesting keyboard-interactive.#
 100         # The hack to workaround this is to attempt a reconnect, which will   #
 101         # receive the right banner, and authentication can proceed. See the   #
 102         # following for more info:                                            #
 103         # https://github.com/paramiko/paramiko/issues/432                     #
 104         # https://github.com/paramiko/paramiko/pull/438                       #
 105         #######################################################################
 106
 107         try:
 108             ssh.connect(host, port=22, username=user, pkey=pkey)
 109         except paramiko.ssh_exception.SSHException as e:
 110             if 'Error reading SSH protocol banner' == str(e):
 111                 # Once more, with feeling
 112                 ssh.connect(host, port=22, username=user, pkey=pkey)
 113             else:
 114                 # Reraise the original exception
 115                 raise e
 116
 117         return ssh
 118
 119     def _run_command(self, ssh, cmd, pty=True):
 120         """Run a command remotely via SSH.
 121
 122         :param object ssh: The SSHClient
 123         :param str cmd: The command to execute
 124         :param list cmd: The `shlex.split` command to execute
 125         :param bool pty: Whether to allocate a pty
 126
 127         :return: tuple: The stdout and stderr of the command execution
 128         :raises: :class:`CalledProcessError` if the command fails
 129         """
 130
 131         if isinstance(cmd, str):
 132             cmd = shlex.split(cmd)
 133
 134         if type(cmd) is not list:
 135             cmd = [cmd]
 136
 137         cmds = ' '.join(cmd)
 138         stdin, stdout, stderr = ssh.exec_command(cmds, get_pty=pty)
 139         retcode = stdout.channel.recv_exit_status()
 140
 141         if retcode > 0:
 142             output = stderr.read().strip()
 143             raise CalledProcessError(returncode=retcode, cmd=cmd,
 144                                      output=output)
 145         return (
 146             stdout.read().decode('utf-8').strip(),
 147             stderr.read().decode('utf-8').strip()
 148         )
 149
 150     def _init_ubuntu_user(self):
 151         """Initialize the ubuntu user.
 152
 153         :return: bool: If the initialization was successful
 154         :raises: :class:`paramiko.ssh_exception.AuthenticationException`
 155             if the authentication fails
 156         """
 157
 158         # TODO: Test this on an image without the ubuntu user setup.
 159
 160         auth_user = self.user
 161         ssh = None
 162         try:
 163             # Run w/o allocating a pty, so we fail if sudo prompts for a passwd
 164             ssh = self._get_ssh_client(
 165                 self.host,
 166                 "ubuntu",
 167                 self.private_key_path,
 168             )
 169
 170             stdout, stderr = self._run_command(ssh, "sudo -n true", pty=False)
 171         except paramiko.ssh_exception.AuthenticationException as e:
 172             raise e
 173         else:
 174             auth_user = "ubuntu"
 175         finally:
 176             if ssh:
 177                 ssh.close()
 178
 179         # if the above fails, run the init script as the authenticated user
 180
 181         # Infer the public key
 182         public_key = None
 183         public_key_path = "{}.pub".format(self.private_key_path)
 184
 185         if not os.path.exists(public_key_path):
 186             raise FileNotFoundError(
 187                 "Public key '{}' doesn't exist.".format(public_key_path)
 188             )
 189
 190         with open(public_key_path, "r") as f:
 191             public_key = f.readline()
 192
 193         script = INITIALIZE_UBUNTU_SCRIPT.format(public_key)
 194
 195         try:
 196             ssh = self._get_ssh_client(
 197                 self.host,
 198                 auth_user,
 199                 self.private_key_path,
 200             )
 201
 202             self._run_command(
 203                 ssh,
 204                 ["sudo", "/bin/bash -c " + shlex.quote(script)],
 205                 pty=True
 206             )
 207         except paramiko.ssh_exception.AuthenticationException as e:
 208             raise e
 209         finally:
 210             ssh.close()
 211
 212         return True
 213
 214     def _detect_hardware_and_os(self, ssh):
 215         """Detect the target hardware capabilities and OS series.
 216
 217         :param object ssh: The SSHClient
 218         :return: str: A raw string containing OS and hardware information.
 219         """
 220
 221         info = {
 222             'series': '',
 223             'arch': '',
 224             'cpu-cores': '',
 225             'mem': '',
 226         }
 227
 228         stdout, stderr = self._run_command(
 229             ssh,
 230             ["sudo", "/bin/bash -c " + shlex.quote(DETECTION_SCRIPT)],
 231             pty=True,
 232         )
 233
 234         lines = stdout.split("\n")
 235         info['series'] = lines[0].strip()
 236         info['arch'] = normalize_arch(lines[1].strip())
 237
 238         memKb = re.split(r'\s+', lines[2])[1]
 239
 240         # Convert megabytes -> kilobytes
 241         info['mem'] = round(int(memKb) / 1024)
 242
 243         # Detect available CPUs
 244         recorded = {}
 245         for line in lines[3:]:
 246             physical_id = ""
 247             print(line)
 248
 249             if line.find("physical id") == 0:
 250                 physical_id = line.split(":")[1].strip()
 251             elif line.find("cpu cores") == 0:
 252                 cores = line.split(":")[1].strip()
 253
 254                 if physical_id not in recorded.keys():
 255                     info['cpu-cores'] += cores
 256                     recorded[physical_id] = True
 257
 258         return info
 259
 260     def provision_machine(self):
 261         """Perform the initial provisioning of the target machine.
 262
 263         :return: bool: The client.AddMachineParams
 264         :raises: :class:`paramiko.ssh_exception.AuthenticationException`
 265             if the upload fails
 266         """
 267         params = client.AddMachineParams()
 268
 269         if self._init_ubuntu_user():
 270             try:
 271
 272                 ssh = self._get_ssh_client(
 273                     self.host,
 274                     self.user,
 275                     self.private_key_path
 276                 )
 277
 278                 hw = self._detect_hardware_and_os(ssh)
 279                 params.series = hw['series']
 280                 params.instance_id = "manual:{}".format(self.host)
 281                 params.nonce = "manual:{}:{}".format(
 282                     self.host,
 283                     str(uuid.uuid4()),  # a nop for Juju w/manual machines
 284                 )
 285                 params.hardware_characteristics = {
 286                     'arch': hw['arch'],
 287                     'mem': int(hw['mem']),
 288                     'cpu-cores': int(hw['cpu-cores']),
 289                 }
 290                 params.addresses = [{
 291                     'value': self.host,
 292                     'type': 'ipv4',
 293                     'scope': 'public',
 294                 }]
 295
 296             except paramiko.ssh_exception.AuthenticationException as e:
 297                 raise e
 298             finally:
 299                 ssh.close()
 300
 301         return params
 302
 303     async def install_agent(self, connection, nonce, machine_id):
 304         """
 305         :param object connection: Connection to Juju API
 306         :param str nonce: The nonce machine specification
 307         :param str machine_id: The id assigned to the machine
 308
 309         :return: bool: If the initialization was successful
 310         """
 311
 312         # The path where the Juju agent should be installed.
 313         data_dir = "/var/lib/juju"
 314
 315         # Disabling this prevents `apt-get update` from running initially, so
 316         # charms will fail to deploy
 317         disable_package_commands = False
 318
 319         client_facade = client.ClientFacade.from_connection(connection)
 320         results = await client_facade.ProvisioningScript(
 321             data_dir,
 322             disable_package_commands,
 323             machine_id,
 324             nonce,
 325         )
 326
 327         self._run_configure_script(results.script)
 328
 329     def _run_configure_script(self, script):
 330         """Run the script to install the Juju agent on the target machine.
 331
 332         :param str script: The script returned by the ProvisioningScript API
 333         :raises: :class:`paramiko.ssh_exception.AuthenticationException`
 334             if the upload fails
 335         """
 336
 337         _, tmpFile = tempfile.mkstemp()
 338         with open(tmpFile, 'w') as f:
 339             f.write(script)
 340
 341         try:
 342             # get ssh client
 343             ssh = self._get_ssh_client(
 344                 self.host,
 345                 "ubuntu",
 346                 self.private_key_path,
 347             )
 348
 349             # copy the local copy of the script to the remote machine
 350             sftp = paramiko.SFTPClient.from_transport(ssh.get_transport())
 351             sftp.put(
 352                 tmpFile,
 353                 tmpFile,
 354             )
 355
 356             # run the provisioning script
 357             stdout, stderr = self._run_command(
 358                 ssh,
 359                 "sudo /bin/bash {}".format(tmpFile),
 360             )
 361
 362         except paramiko.ssh_exception.AuthenticationException as e:
 363             raise e
 364         finally:
 365             os.remove(tmpFile)
 366             ssh.close()