build.py

#!/usr/bin/env python3
"""Build script for creating AWS lambda function and layer artifacts.

This script generally lives in a build pipeline and will create zipfile archive
artifacts that are suitable for deployment to a lambda function or layer.
"""

from __future__ import annotations

import argparse
import os
import pprint
import shutil
import subprocess
import warnings
import zipfile
from dataclasses import dataclass
from enum import Enum
from functools import cached_property
from pathlib import Path
from typing import Dict, List, Optional, Set
from zipfile import ZipFile


class Language(Enum):
    """Enumerated types of supported languages."""

    PYTHON = "py"
    TYPESCRIPT = "ts"
    JAVASCRIPT = "js"
    GOLANG = "go"


LAMBDA_SIZE_LIMIT_BYTES = 50000000

DEFAULT_ALLOWED_FILE_TYPES = set(
    [
        ".9",
        ".APACHE",
        ".BSD",
        ".PL",
        ".PSF",
        ".R",
        ".TAG",
        ".bat",
        ".c",
        ".cc",
        ".cfg",
        ".csh",
        ".css",
        ".exe",
        ".fish",
        ".gemspec",
        ".go",
        ".gz",
        ".h",
        ".hash",
        ".html",
        ".html",
        ".ini",
        ".js",
        ".json",
        ".lock",
        ".md",
        ".mod",
        ".npmignore",
        ".nu",
        ".pem",
        ".pem",
        ".png",
        ".properties",
        ".ps1",
        ".pth",
        ".pump",
        ".pxd",
        ".pxi",
        ".py",
        ".pyc",
        ".pyi",
        ".pyx",
        ".renv",
        ".rockspec",
        ".rs",
        ".rst",
        ".scss",
        ".sh",
        ".so",
        ".test",
        ".tmpl",
        ".toml",
        ".txt",
        ".typed",
        ".virtualenv",
        ".whl",
        ".xml",
        ".xsd",
        ".xslt",
        ".yaml",
        ".yml",
        "Makefile",
    ]
)

REPO_ROOT = Path(os.path.dirname((os.path.dirname(os.path.realpath(__file__)))))


class BuildFailureError(RuntimeError):
    """Subclass of terminal build failure error."""

    def __init__(self, command: str, perror: bytes):
        """Wrap command and error output."""
        error_string = perror.rstrip().decode("utf-8")
        super().__init__(
            f"Failed to execute build for command '{command}' due to this error:"
            f" '{error_string}'"
        )
        self.command: list = command


@dataclass
class LambdaArtifact:
    """Base Lambda artifact object, shared logic for all python artifacts.

    This is a python specific base object.  It will work for other languages,
    but is_skippable methods should be altered.
    """

    language: Language
    src_root: Path
    artifact_name: str
    allowed_file_types: Optional[Set[str]] = None
    ignored_file_types: Optional[Set[str]] = None

    def __post_init__(self) -> None:
        """Set mutable defaults, allow for additional types."""
        if self.allowed_file_types:
            self.allowed_file_types = DEFAULT_ALLOWED_FILE_TYPES | set(
                self.allowed_file_types
            )
        else:
            self.allowed_file_types = DEFAULT_ALLOWED_FILE_TYPES

        if not isinstance(self.language, Language):
            raise RuntimeError(
                "Attempt to build an artifact in an unsupported language"
            )

    @cached_property
    def repo_root(self) -> Path:
        """Syntactic sugar for repo root path."""
        return Path(os.path.dirname((os.path.dirname(os.path.realpath(__file__)))))

    def execute(self, _cmd: str) -> tuple:
        """
        Shell out and run a bash command.

        Purpose  : To execute a command and return exit statuns
        Argument : cmd - command to execute
        Return   : result, exit_code
        """
        warnings.warn(
            "Shelling out to bash, please avoid this if possible",
            DeprecationWarning,
            stacklevel=2,
        )
        process = subprocess.Popen(
            _cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE
        )
        (result, error) = process.communicate()
        r_c = process.wait()

        if r_c != 0:
            raise BuildFailureError(_cmd, error)

        return result.rstrip().decode("utf-8"), error.rstrip().decode("utf-8")

    def is_root_skippable(self, root) -> bool:
        """Determine if we can ignore the directory when making the artifact.

        Given a 'root' from os.walk, can we skip the whole directory?
        """
        if (
            # ".venv" in root
            "test" == root[-4:]
            or "/test/" in root
            or "__pycache__" in root
            or ".pytest_cache" in root
            # or "/fixtures/" in root
        ):
            return True

    def is_file_skippable(self, relative_path) -> bool:
        """Determine if we can ignore the file when making the artifact.

        Given a relative path from os.walk, can we skip the file?
        Paths are relative the the zipfile, ie python/boto3/resources/model.py
        """
        filename, file_extension = os.path.splitext(relative_path)
        if filename == "handler":
            return False

        if not file_extension:
            # It's an empty directory
            return True

        if (
            "test_" == filename[:5]
            or "conftest" == filename
            or "pytest" in filename
            or "." == filename[0]
            or "conftest.py" == filename
        ):
            return True

        if file_extension not in self.allowed_file_types:
            return True

        return False

    def create(self, path, ziph) -> None:
        """Zip all the files in a directory, recursively.

        Possible future optimization:
        - https://stackoverflow.com/questions/19859840/excluding-directories-in-os-walk
        """
        allowed_file_types = self.allowed_file_types

        self.ignored_file_types = set()
        self.ignored_filenames = set()
        self.ignored_roots = set()

        # We don't use dirs here, it's returned by walk so we set it, but we
        # don't use it
        for root, dirs, filenames in os.walk(path):
            if self.is_root_skippable(root):
                self.ignored_roots.add(root)  # but not really
                continue

            for relative_path in filenames:
                filename, file_extension = os.path.splitext(relative_path)

                if self.is_file_skippable(relative_path):
                    self.ignored_filenames.add(root)
                    continue

                if file_extension not in allowed_file_types:
                    self.ignored_file_types.add(file_extension)

                ziph.write(
                    os.path.join(root, relative_path),
                    os.path.relpath(
                        os.path.join(root, relative_path), os.path.join(path, ".")
                    ),
                )

    def compile(self) -> None:
        """Abstract method, do nothing if not extended."""
        pass

    def is_python(self) -> bool:
        """Return if this artifact is python."""
        return self.language == Language.PYTHON

    def is_javascript(self) -> bool:
        """Return if this artifact is javascript."""
        return self.language == Language.JAVASCRIPT

    def is_typescript(self) -> bool:
        """Return if this artifact is typescript."""
        return self.language == Language.TYPESCRIPT

    def is_golang(self) -> bool:
        """Return if this artifact is golang."""
        return self.language == Language.GOLANG

    @property
    def artifact_root(self) -> Path:
        """Syntactic sugar for artifact root path."""
        return Path(REPO_ROOT.joinpath(self.src_root))

    def zipdir(self, path, ziph) -> LambdaArtifact:
        """Zip all the files in a directory, recursively.

        Possible future optimization:
        - https://stackoverflow.com/questions/19859840/excluding-directories-in-os-walk
        """
        self.create(path, ziph)

        # TODO: get working again
        # self.output(f"Skipped file roots {pprint.pformat(artifact.ignored_roots)}")
        # self.output(f"Skipped file types {pprint.pformat(artifact.ignored_file_types)}")
        # self.output(f"Skipped filenames {pprint.pformat(artifact.ignored_filenames)}")

        return self

    def create_archive(self, src_path) -> tuple:
        """Create an achival aritfact.

        src_path is path to the directory to zip up
        example: PosixPath('/home/ahonnecke/src/rsu-scrapers/build/tmp/BatchQueryDeviceLambda')
        zip_path is ROOT/build/<artifact_name>.zip
        example: PosixPath('/home/ahonnecke/src/rsu-scrapers/build/BatchQueryDeviceLambda.zip')
        """
        # Add artifact to the list of expected artifacts for later verification
        os.chdir(REPO_ROOT)
        build_path = REPO_ROOT.joinpath("build")
        zip_path = build_path.joinpath(f"{self.artifact_name}.zip")
        full_path = self.repo_root.joinpath(src_path)

        with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
            self.zipdir(full_path, zipf)

    @property
    def filename(self) -> str:
        """Syntactic sugar for filename of artifact."""
        return f"{self.artifact_name}.zip"

    def validate_zip(self) -> None:
        """Validate the artifact zip file."""
        self.validate_zip_size()
        self.validate_zip_contents()

    @property
    def zipped_archive(self) -> ZipFile:
        """Syntactic sugar for extracting the zipfile artifact."""
        try:
            return zipfile.ZipFile(self.filename)
        except Exception as e:
            raise RuntimeError(f"The {self.filename} artifact is malformed.") from e

    def validate_zip_size(self) -> None:
        """Extract zipfile and ensure it's not too big."""
        os.chdir(REPO_ROOT.joinpath("build"))

        if not os.path.isfile(self.filename):
            assert False, f"File {self.filename} does not exist"

        file_size = os.path.getsize(self.filename)
        human_readable = f"which is {file_size/float(1<<20):,.0f} MB"

        if file_size > LAMBDA_SIZE_LIMIT_BYTES:
            raise RuntimeError(
                f"Zipfile {self.filename} is too damn big ({file_size}) {human_readable}"
            )

        print(f"Zipfile {self.filename} is OK ({file_size}) {human_readable}")

        assert self.zipped_archive.namelist()
        assert self.zipped_archive.infolist()

    def validate_zip_contents(self):
        self.suspicious_files = []

        for singlefile in self.zipped_archive.namelist():
            filename, file_extension = os.path.splitext(singlefile)
            if file_extension not in self.allowed_file_types and filename != "handler":
                self.suspicious_files.append(file_extension)

        mylen = len(self.zipped_archive.namelist())
        assert mylen, f"Zipfile {self.filename} is empty"

        if self.suspicious_files:
            warnings.warn(f"{len(self.suspicious_files)} unexpected file types found")

            pprint.pformat(self.suspicious_files, indent=4)

        return zip

    @property
    def build_space(self) -> Path:
        """Syntactic sugar for build space path."""
        return Path(f"{REPO_ROOT}/build/tmp/{self.artifact_name}")


class LambdaFunctionArtifact(LambdaArtifact):
    """Object for creation of artifact for a lambda function.

    Contains logic that's specific to functions.
    """

    def validate_function() -> None:
        os.chdir(REPO_ROOT.joinpath("build"))

    def create_archive(self) -> Dict:
        return super().create_archive(self.dest_root)

    @property
    def dest_root(self) -> Path:
        """Syntactic sugar for build step destination."""
        if self.is_typescript():
            return self.artifact_root.parent.joinpath("dist/")
        else:
            return self.src_root

    def npm_transpile(self) -> tuple:
        """Shell out and install and build, transpiling to javascript."""
        self.execute("npm install")
        return self.execute("npm run build")

    def go_compile(self) -> bool:
        """Perform the steps to compile golang into a binary."""
        CWD = os.getcwd()
        os.chdir(self.artifact_root)

        if os.getenv("GITHUB_ACTIONS"):
            print("Do CI specific steps")

        src = REPO_ROOT.joinpath(self.src_root)
        dest = self.build_space.joinpath("handler")

        try:
            os.remove(dest)
        except FileNotFoundError:
            pass

        self.execute(f"go build -o {dest} {src}")
        self.src_root = self.build_space

        os.chdir(CWD)
        return True

    def compile(self) -> bool:
        """Compile, or not as is needed by the language."""
        if self.is_python():
            pass
        elif self.is_javascript():
            pass
        elif self.is_typescript():
            os.chdir(self.artifact_root.parent)
            self.npm_transpile()
        elif self.is_golang():
            return self.go_compile()
        return True


class LambdaLayerArtifact(LambdaArtifact):
    """Object for creation of artifact for a lambda layer.

    Contains logic that's specific to layers.
    """

    def create_archive(self) -> Dict:
        """Build the zipfile artifact."""
        return super().create_archive(self.build_space)

    def compile(self) -> None:
        """Compile the layer artifact."""
        root = REPO_ROOT.joinpath(self.src_root)
        self.build_space.mkdir(parents=True, exist_ok=True)

        os.chdir(root)

        if self.is_python():
            # TODO: get pipenv stuff from device notifcation build
            print("Compiling py layer")
        elif self.is_javascript():
            print("Compiling Js layer")

            node_js = self.build_space.joinpath("nodejs")
            node_js.mkdir(parents=True, exist_ok=True)

            os.chdir(self.build_space)
            pkg = REPO_ROOT.joinpath(self.src_root).joinpath("package.json")
            shutil.copyfile(pkg, node_js.joinpath("package.json"))

            os.chdir(node_js)
            self.npm_install_layer()

        elif self.is_golang():
            print("Compiling GO layer")

    def npm_install_layer(self) -> tuple:
        """Shell out and install prod dependencies."""
        return self.execute("npm install --omit=dev")


@dataclass
class ZippityDooDah:
    """Combined sets of functions and layers to produce artifacts for.

    Empty lists are allowed, but a list is required.
    """

    functions: List[LambdaFunctionArtifact]
    layers: List[LambdaLayerArtifact]
    quiet: bool = False

    artifacts = []

    @cached_property
    def root(self) -> Path:
        """Syntactic sugar for repo root."""
        return Path(os.path.dirname((os.path.dirname(os.path.realpath(__file__)))))

    @staticmethod
    def get_args():
        parser = argparse.ArgumentParser(description="Build repository artifacts")

        parser.add_argument(
            "-n",
            "--no-build",
            action="store_true",
            default=False,
            help="Skip the build step, helpful for testing",
        )

        parser.add_argument(
            "-q",
            "--quiet",
            action="store_true",
            default=False,
            help="Suppress output.",
        )

        return parser.parse_args()

    def zipdir(self, path, ziph, artifact: LambdaArtifact):
        """Zip all the files in a directory, recursively.

        Possible future optimization:
        - https://stackoverflow.com/questions/19859840/excluding-directories-in-os-walk
        """
        artifact.create(path, ziph)

        self.output(f"Skipped file roots {pprint.pformat(artifact.ignored_roots)}")
        self.output(f"Skipped file types {pprint.pformat(artifact.ignored_file_types)}")
        self.output(f"Skipped filenames {pprint.pformat(artifact.ignored_filenames)}")

        return artifact

    def output(self, message: str, _=None) -> None:
        if not self.quiet:
            print(message)

    def section(self, message: str, _=None) -> None:
        self.output("========================================================")
        self.output(message)
        self.output("========================================================")

    def prepare_buildspace(self) -> None:
        self.section(f"Preparing workspace {self.root}")

        build_dir = REPO_ROOT.joinpath("build")
        build_dir.mkdir(parents=True, exist_ok=True)

        for file in self.artifacts:
            try:
                os.remove(build_dir.joinpath(file.get("filename")))
            except FileNotFoundError:
                pass
    
    def cleanup_buildspace(self) -> None:
        self.section(f"Cleaning up workspace {self.root}")
        
        build_dir = REPO_ROOT.joinpath("build")

        try:
            shutil.rmtree(build_dir.joinpath("tmp"))
            os.remove(build_dir.joinpath(".gitignore"))
        except:
            pass

    def build(self) -> None:
        self.prepare_buildspace()

        self.section("Starting function artifacts")
        for _func in self.functions:
            _func.compile()
            _func.create_archive()
            _func.validate_zip()

        self.section("Starting layer artifacts")
        for _layer in self.layers:
            _layer.compile()
            _layer.create_archive()
            _layer.validate_zip()

        self.cleanup_buildspace()
        self.section("Validation complete")


def main() -> None:
    """Run main program body.

    This contains all the relevant information for
    building functions and layers, if you are reusing ZippityDooDah, this is the
    only location where changes are needeed
    """
    args = ZippityDooDah.get_args()
    zippy = ZippityDooDah(
        quiet=args.quiet,
        functions=[
            LambdaFunctionArtifact(
                src_root="device-data-collector/device-query-transform/src/",
                artifact_name="DeviceQueryTransformLambda",
                language=Language.JAVASCRIPT,
            ),
            LambdaFunctionArtifact(
                src_root="device-data-collector/enqueue-rsus-for-data-collection/src/",
                artifact_name="EnqueueLambda",
                language=Language.TYPESCRIPT,
            ),
            LambdaFunctionArtifact(
                src_root="device-data-collector/batch-query-device",
                artifact_name="BatchQueryDeviceLambda",
                language=Language.GOLANG,
            ),
        ],
        layers=[
            LambdaLayerArtifact(
                src_root="device-data-collector/enqueue-rsus-for-data-collection",
                artifact_name="EnqueueLayer",
                language=Language.JAVASCRIPT,
            ),
            LambdaLayerArtifact(
                src_root="device-data-collector/device-query-transform/",
                artifact_name="DeviceQueryTransformLayer",
                language=Language.JAVASCRIPT,
            ),
        ],
    )
    zippy.build()


if __name__ == "__main__":
    main()