artifacts.py

# Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0
"""Define classes for interacting with CI artifacts in s3."""

import os
import platform
import tempfile
from shutil import copyfile
from enum import Enum
from stat import S_IREAD, S_IWRITE
from pathlib import Path
import boto3
import botocore.client
from framework.defs import DEFAULT_TEST_SESSION_ROOT_PATH
from framework.utils import compare_versions
from host_tools.snapshot_helper import merge_memory_bitmaps


ARTIFACTS_LOCAL_ROOT = f"{DEFAULT_TEST_SESSION_ROOT_PATH}/ci-artifacts"


class ArtifactType(Enum):
    """Supported artifact types."""

    MICROVM = "microvm"
    KERNEL = "kernel"
    DISK = "disk"
    SSH_KEY = "ssh_key"
    MEM = "mem"
    VMSTATE = "vmstate"
    MISC = "misc"
    SNAPSHOT = "snapshot"
    FC = "firecracker"
    JAILER = "jailer"


class Artifact:
    """A generic read-only artifact manipulation class."""

    LOCAL_ARTIFACT_DIR = f"{DEFAULT_TEST_SESSION_ROOT_PATH}/local-artifacts"

    def __init__(self,
                 bucket,
                 key,
                 artifact_type=ArtifactType.MISC,
                 local_folder=None,
                 is_copy=False):
        """Initialize bucket, key and type."""
        self._bucket = bucket
        self._key = key
        self._local_folder = local_folder
        self._type = artifact_type
        self._is_copy = is_copy

    @property
    def type(self):
        """Return the artifact type."""
        return self._type

    @property
    def key(self):
        """Return the artifact key."""
        return self._key

    @property
    def bucket(self):
        """Return the artifact bucket."""
        return self._bucket

    def name(self):
        """Return the artifact name."""
        return Path(self.key).name

    def base_name(self):
        """Return the base name (without extension)."""
        return Path(self.key).stem

    def local_dir(self):
        """Return the directory containing the downloaded artifact."""
        assert self._local_folder is not None
        return "{}/{}/{}".format(
            self._local_folder,
            self.type.value,
            platform.machine(),
        )

    def download(self, target_folder=ARTIFACTS_LOCAL_ROOT, force=False):
        """Save the artifact in the folder specified target_path."""
        assert self.bucket is not None
        self._local_folder = target_folder
        Path(self.local_dir()).mkdir(parents=True, exist_ok=True)
        if force or not os.path.exists(self.local_path()):
            self._bucket.download_file(self._key, self.local_path())
            # Artifacts are read only by design.
            os.chmod(self.local_path(), S_IREAD)

    def local_path(self):
        """Return the local path where the file was downloaded."""
        # The file path format: <target_folder>/<type>/<platform>/<name>
        return "{}/{}".format(
            self.local_dir(),
            self.name()
        )

    def copy(self, file_name=None):
        """Create a writeable copy of the artifact."""
        assert os.path.exists(self.local_path()), """File {} not found.
        call download() first.""".format(self.local_path())

        # The file path for this artifact copy.
        new_dir = "{}/{}/{}".format(
            Artifact.LOCAL_ARTIFACT_DIR,
            self.type.value,
            platform.machine()
        )

        if file_name is None:
            # The temp file suffix is the artifact type.
            suffix = "-{}".format(self.type.value)
            # The key for the new artifact is the full path to the file.
            new_key = tempfile.mktemp(dir=new_dir, suffix=suffix)
        else:
            # Caller specified new name.
            new_key = os.path.join(new_dir, file_name)

        # Create directories if needed.
        Path(new_dir).mkdir(parents=True, exist_ok=True)
        # Copy to local artifact.
        copyfile(self.local_path(), new_key)
        # Make it writeable.
        os.chmod(new_key, S_IREAD | S_IWRITE)
        # Local folder of the new artifact.
        local_folder = Artifact.LOCAL_ARTIFACT_DIR
        # Calls to download() on the new Artifact are guarded by a
        # bucket assert.
        return Artifact(None, new_key,
                        artifact_type=self.type,
                        local_folder=local_folder, is_copy=True)

    def __del__(self):
        """Teardown the object."""
        if self._is_copy and os.path.exists(self._key):
            os.remove(self._key)


class SnapshotArtifact:
    """Manages snapshot S3 artifact objects."""

    def __init__(self,
                 bucket,
                 key,
                 artifact_type=ArtifactType.SNAPSHOT):
        """Initialize bucket, key and type."""
        self._bucket = bucket
        self._type = artifact_type
        self._key = key

        self._mem = Artifact(self._bucket, "{}vm.mem".format(key),
                             artifact_type=ArtifactType.MISC)
        self._vmstate = Artifact(self._bucket, "{}vm.vmstate".format(key),
                                 artifact_type=ArtifactType.MISC)
        self._ssh_key = Artifact(self._bucket, "{}ssh_key".format(key),
                                 artifact_type=ArtifactType.SSH_KEY)
        self._disks = []

        disk_prefix = "{}disk".format(key)
        snaphot_disks = self._bucket.objects.filter(Prefix=disk_prefix)

        for disk in snaphot_disks:
            artifact = Artifact(self._bucket, disk.key,
                                artifact_type=ArtifactType.DISK)
            self._disks.append(artifact)

        # Get the name of the snapshot folder.
        snapshot_name = self.name
        self._local_folder = os.path.join(ARTIFACTS_LOCAL_ROOT,
                                          self.type.value,
                                          snapshot_name)

    @property
    def type(self):
        """Return the artifact type."""
        return self._type

    @property
    def key(self):
        """Return the artifact key."""
        return self._key

    @property
    def mem(self):
        """Return the memory artifact."""
        return self._mem

    @property
    def vmstate(self):
        """Return the vmstate artifact."""
        return self._vmstate

    @property
    def ssh_key(self):
        """Return the vmstate artifact."""
        return self._ssh_key

    @property
    def disks(self):
        """Return the disk artifacts."""
        return self._disks

    @property
    def name(self):
        """Return the name of the artifact."""
        return self._key.strip('/').split('/')[-1]

    def download(self):
        """Download artifacts and return a Snapshot object."""
        self.mem.download(self._local_folder)
        self.vmstate.download(self._local_folder)
        # SSH key is not needed by microvm, it is needed only by
        # test functions.
        self.ssh_key.download(self._local_folder)

        for disk in self.disks:
            disk.download(self._local_folder)
            os.chmod(disk.local_path(), 0o700)

    def copy(self, vm_root_folder):
        """Copy artifacts and return a Snapshot object."""
        assert self._local_folder is not None

        dst_mem_path = os.path.join(vm_root_folder, self.mem.name())
        dst_state_file = os.path.join(vm_root_folder, self.vmstate.name())
        dst_ssh_key = os.path.join(vm_root_folder, self.ssh_key.name())

        # Copy mem, state & ssh_key.
        copyfile(self.mem.local_path(), dst_mem_path)
        copyfile(self.vmstate.local_path(), dst_state_file)
        copyfile(self.ssh_key.local_path(), dst_ssh_key)
        # Set proper permissions for ssh key.
        os.chmod(dst_ssh_key, 0o400)

        disk_paths = []
        for disk in self.disks:
            dst_disk_path = os.path.join(vm_root_folder, disk.name())
            copyfile(disk.local_path(), dst_disk_path)
            disk_paths.append(dst_disk_path)

        return Snapshot(dst_mem_path,
                        dst_state_file,
                        disks=disk_paths,
                        net_ifaces=None,
                        ssh_key=dst_ssh_key)


class DiskArtifact(Artifact):
    """Provides access to associated ssh key artifact."""

    def ssh_key(self):
        """Return a ssh key artifact."""
        # Replace extension.
        key_file_path = str(Path(self.key).with_suffix('.id_rsa'))
        return Artifact(self.bucket,
                        key_file_path,
                        artifact_type=ArtifactType.SSH_KEY)


class FirecrackerArtifact(Artifact):
    """Provides access to associated jailer artifact."""

    def jailer(self):
        """Return a jailer binary artifact."""
        # Jailer and FC binaries have different extensions and share
        # file name when stored in S3:
        # Firecracker binary: v0.22.firecrcker
        # Jailer binary: v0.23.0.jailer
        jailer_path = str(Path(self.key).with_suffix('.jailer'))
        return Artifact(self.bucket,
                        jailer_path,
                        artifact_type=ArtifactType.JAILER)

    @property
    def version(self):
        """Return the artifact's version: `X.Y.Z`."""
        # Get the filename, remove the extension and trim the leading 'v'.
        return os.path.splitext(os.path.basename(self.key))[0][1:]


class ArtifactCollection:
    """Provides easy access to different artifact types."""

    MICROVM_CONFIG_EXTENSION = ".json"
    MICROVM_KERNEL_EXTENSION = ".bin"
    MICROVM_DISK_EXTENSION = ".ext4"
    MICROVM_VMSTATE_EXTENSION = ".vmstate"
    MICROVM_MEM_EXTENSION = ".mem"
    FC_EXTENSION = ".firecracker"
    JAILER_EXTENSION = ".jailer"

    PLATFORM = platform.machine()

    # S3 bucket structure.
    ARTIFACTS_ROOT = 'ci-artifacts'
    ARTIFACTS_DISKS = '/disks/' + PLATFORM + "/"
    ARTIFACTS_KERNELS = '/kernels/' + PLATFORM + "/"
    ARTIFACTS_MICROVMS = '/microvms/'
    ARTIFACTS_SNAPSHOTS = '/snapshots/' + PLATFORM + "/"
    ARTIFACTS_BINARIES = '/binaries/' + PLATFORM + "/"

    def __init__(
        self,
        bucket
    ):
        """Initialize S3 client."""
        config = botocore.client.Config(signature_version=botocore.UNSIGNED)
        # pylint: disable=E1101
        # fixes "E1101: Instance of '' has no 'Bucket' member (no-member)""
        self.bucket = boto3.resource('s3', config=config).Bucket(bucket)

    def _fetch_artifacts(self,
                         artifact_dir,
                         artifact_ext,
                         artifact_type,
                         artifact_class,
                         keyword=None):
        artifacts = []
        prefix = ArtifactCollection.ARTIFACTS_ROOT + artifact_dir
        files = self.bucket.objects.filter(Prefix=prefix)
        for file in files:
            if (
                # Filter by extensions.
                file.key.endswith(artifact_ext)
                # Filter by userprovided keyword if any.
                and (keyword is None or keyword in file.key)
            ):
                artifacts.append(artifact_class(self.bucket,
                                                file.key,
                                                artifact_type=artifact_type))
        return artifacts

    def snapshots(self, keyword=None):
        """Return snapshot artifacts for the current arch."""
        # Snapshot artifacts are special since they need to contain
        # a variable number of files: mem, state, disks, ssh key.
        # To simplify the way we retrieve and store snapshot artifacts
        # we are going to group all snapshot file in a folder and the
        # "keyword" parameter will filter this folder name.
        #
        # Naming convention for files within the snapshot below.
        # Snapshot folder /ci-artifacts/snapshots/x86_64/fc_snapshot_v0.22:
        # - vm.mem
        # - vm.vmstate
        # - disk0 <---- this is the root disk
        # - disk1
        # - diskN
        # - ssh_key

        artifacts = []
        prefix = ArtifactCollection.ARTIFACTS_ROOT
        prefix += ArtifactCollection.ARTIFACTS_SNAPSHOTS
        snaphot_dirs = self.bucket.objects.filter(Prefix=prefix)
        for snapshot_dir in snaphot_dirs:
            key = snapshot_dir.key
            # Filter out the snapshot artifacts root folder.
            # Select only files with specified keyword.
            if (key[-1] == "/" and key != prefix and
                    (keyword is None or keyword in snapshot_dir.key)):
                artifact_type = ArtifactType.SNAPSHOT
                artifacts.append(SnapshotArtifact(self.bucket,
                                                  key,
                                                  artifact_type=artifact_type))

        return artifacts

    def microvms(self, keyword=None):
        """Return microvms artifacts for the current arch."""
        return self._fetch_artifacts(
            ArtifactCollection.ARTIFACTS_MICROVMS,
            ArtifactCollection.MICROVM_CONFIG_EXTENSION,
            ArtifactType.MICROVM,
            Artifact,
            keyword=keyword
        )

    def firecrackers(self, keyword=None, older_than=None):
        """Return fc/jailer artifacts for the current arch."""
        firecrackers = self._fetch_artifacts(
            ArtifactCollection.ARTIFACTS_BINARIES,
            ArtifactCollection.FC_EXTENSION,
            ArtifactType.FC,
            FirecrackerArtifact,
            keyword=keyword
        )

        # Filter out binaries with versions newer than the `older_than` arg.
        if older_than is not None:
            return list(filter(
                lambda fc: compare_versions(fc.version, older_than) <= 0,
                firecrackers
            ))

        return firecrackers

    def firecracker_versions(self, older_than=None):
        """Return fc/jailer artifacts' versions for the current arch."""
        return [fc.base_name()[1:]
                for fc in self.firecrackers(older_than=older_than)]

    def kernels(self, keyword=None):
        """Return kernel artifacts for the current arch."""
        return self._fetch_artifacts(
            ArtifactCollection.ARTIFACTS_KERNELS,
            ArtifactCollection.MICROVM_KERNEL_EXTENSION,
            ArtifactType.KERNEL,
            Artifact,
            keyword=keyword
        )

    def disks(self, keyword=None):
        """Return disk artifacts for the current arch."""
        return self._fetch_artifacts(
            ArtifactCollection.ARTIFACTS_DISKS,
            ArtifactCollection.MICROVM_DISK_EXTENSION,
            ArtifactType.DISK,
            DiskArtifact,
            keyword=keyword
        )


class ArtifactSet:
    """Manages a set of artifacts with the same type."""

    def __init__(self, artifacts):
        """Initialize type and artifact array."""
        self._type = None
        self._artifacts = []
        self.insert(artifacts)

    def insert(self, artifacts):
        """Add artifacts to set."""
        if artifacts is not None and len(artifacts) > 0:
            self._type = self._type or artifacts[0].type
        for artifact in artifacts:
            assert artifact.type == self._type
            self._artifacts.append(artifact)

    @property
    def artifacts(self):
        """Return the artifacts array."""
        return self._artifacts

    def __len__(self):
        """Return the artifacts array len."""
        return len(self._artifacts)


class SnapshotType(Enum):
    """Supported snapshot types."""

    FULL = 0
    DIFF = 1


class Snapshot:
    """Manages Firecracker snapshots."""

    def __init__(self, mem, vmstate, disks, net_ifaces, ssh_key):
        """Initialize mem, vmstate, disks, key."""
        assert mem is not None
        assert vmstate is not None
        assert disks is not None
        assert ssh_key is not None
        self._mem = mem
        self._vmstate = vmstate
        self._disks = disks
        self._ssh_key = ssh_key
        self._net_ifaces = net_ifaces

    def rebase_snapshot(self, base):
        """Rebases current incremental snapshot onto a specified base layer."""
        merge_memory_bitmaps(base.mem, self.mem)
        self._mem = base.mem

    def cleanup(self):
        """Delete the backing files from disk."""
        os.remove(self._mem)
        os.remove(self._vmstate)

    @property
    def mem(self):
        """Return the mem file path."""
        return self._mem

    @property
    def vmstate(self):
        """Return the vmstate file path."""
        return self._vmstate

    @property
    def disks(self):
        """Return the disk file paths."""
        return self._disks

    @property
    def ssh_key(self):
        """Return the ssh key file path."""
        return self._ssh_key

    @property
    def net_ifaces(self):
        """Return the list of net interface configs."""
        return self._net_ifaces


# Default configuration values for network interfaces.
DEFAULT_HOST_IP = "192.168.0.1"
DEFAULT_GUEST_IP = "192.168.0.2"
DEFAULT_TAP_NAME = "tap0"
DEFAULT_DEV_NAME = "eth0"
DEFAULT_NETMASK = 30


class NetIfaceConfig:
    """Defines a network interface configuration."""

    def __init__(self,
                 host_ip=DEFAULT_HOST_IP,
                 guest_ip=DEFAULT_GUEST_IP,
                 tap_name=DEFAULT_TAP_NAME,
                 dev_name=DEFAULT_DEV_NAME,
                 netmask=DEFAULT_NETMASK):
        """Initialize object."""
        self._host_ip = host_ip
        self._guest_ip = guest_ip
        self._tap_name = tap_name
        self._dev_name = dev_name
        self._netmask = netmask

    @property
    def host_ip(self):
        """Return the host IP."""
        return self._host_ip

    @property
    def guest_ip(self):
        """Return the guest IP."""
        return self._guest_ip

    @property
    def tap_name(self):
        """Return the tap device name."""
        return self._tap_name

    @property
    def dev_name(self):
        """Return the guest device name."""
        return self._dev_name

    @property
    def netmask(self):
        """Return the netmask."""
        return self._netmask