recipe_modules/ssh/api.py

# Copyright 2020 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import attr

from recipe_engine import recipe_api
from recipe_engine.config_types import Path

_SSH_CONFIG_TEMPLATE = """
Host *
  CheckHostIP no
  StrictHostKeyChecking no
  ForwardAgent no
  ForwardX11 no
  UserKnownHostsFile /dev/null
  User fuchsia
  IdentitiesOnly yes
  IdentityFile {identity}
  ServerAliveInterval 2
  ServerAliveCountMax 5
  ControlMaster auto
  ControlPersist 1m
  ControlPath /tmp/ssh-%r@%h:%p
  ConnectTimeout 5
"""


@attr.s
class SSHFilePaths(object):
    """Required files to setup SSH on FEMU."""

    # Recipe API, required
    _api = attr.ib(type=recipe_api.RecipeApi)

    # Files for SSH
    host_private = attr.ib(type=Path, default=None)
    host_public = attr.ib(type=Path, default=None)

    id_private = attr.ib(type=Path, default=None)
    id_public = attr.ib(type=Path, default=None)

    def _exists(self, p):
        return p and self._api.path.exists(p)

    def _exist(self):
        return all(
            [
                self._exists(self.host_private),
                self._exists(self.host_public),
                self._exists(self.id_private),
                self._exists(self.id_public),
            ]
        )

    def _report_missing(self):
        result = []
        if not self._exists(self.host_private):
            result.append(self.host_private)
        if not self._exists(self.host_public):
            result.append(self.host_public)
        if not self._exists(self.id_private):
            result.append(self.id_private)
        if not self._exists(self.id_public):
            result.append(self.id_public)
        return result


class SSHApi(recipe_api.RecipeApi):
    def __init__(self, *args, **kwargs):
        super(SSHApi, self).__init__(*args, **kwargs)
        self._ssh_paths = None

    def _create_ssh_keys(self, timeout_secs=10 * 60):
        """Generate private, public key-pairs for Host side and Device side ssh keys."""
        self.m.file.ensure_directory("init ssh cache", self.ssh_cache_root)
        if not self._ssh_paths:
            self._ssh_paths = SSHFilePaths(
                api=self.m,
                host_private=self.ssh_cache_root.join("ssh_host_key"),
                host_public=self.ssh_cache_root.join("ssh_host_key.pub"),
                id_private=self.ssh_cache_root.join("id_ed25519"),
                id_public=self.ssh_cache_root.join("id_ed25519.pub"),
            )

        if not self.m.file.listdir(
            name="check ssh cache content", source=self.ssh_cache_root, test_data=()
        ):
            self.m.step(
                "ssh-keygen host",
                [
                    "ssh-keygen",
                    "-t",
                    "ed25519",
                    "-h",
                    "-f",
                    self._ssh_paths.host_private,
                    "-P",
                    "",
                    "-N",
                    "",
                ],
                infra_step=True,
                timeout=timeout_secs,
            )
            self.m.step(
                "ssh-keygen device",
                [
                    "ssh-keygen",
                    "-t",
                    "ed25519",
                    "-f",
                    self._ssh_paths.id_private,
                    "-P",
                    "",
                    "-N",
                    "",
                ],
                infra_step=True,
                timeout=timeout_secs,
            )
        return self._ssh_paths

    def generate_ssh_config(self, private_key_path, dest):
        """Generates and sets the private_key_path in ssh_config file."""
        self.m.file.write_text(
            name="generate ssh_config at %s" % dest,
            dest=dest,
            text_data=_SSH_CONFIG_TEMPLATE.format(identity=private_key_path),
        )

    @property
    def ssh_paths(self):
        """Generate SSH keys.

        Raises:
            StepFailure: When ssh key file paths do not exist.
        """
        self._create_ssh_keys()
        if not self._ssh_paths._exist():
            missing = self._ssh_paths._report_missing()
            ex = self.m.step.StepFailure(
                "SSH paths do not exist. {missing_paths}".format(missing_paths=missing)
            )
            ex.missing_paths = missing
            raise ex
        return self._ssh_paths

    @property
    def ssh_cache_root(self):
        return self.m.buildbucket.builder_cache_path.join("ssh")