majuna/app/terraform/__init__.py

import os
import stat
from typing import Any, Optional, Tuple
from zipfile import ZIP_DEFLATED, ZipFile, ZipInfo

import jinja2


class DeterministicZip:
    """
    Create a zip file deterministically.

    Heavily inspired by https://github.com/bboe/deterministic_zip.
    """

    zipfile: ZipFile

    def __init__(self, filename: str):
        self.zipfile = ZipFile(filename, "w")

    def __enter__(self) -> "DeterministicZip":
        return self

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        self.zipfile.close()

    def add_file(self, path: str, contents: bytes) -> None:
        permission = 0o555 if os.access(path, os.X_OK) else 0o444
        zip_info = ZipInfo()
        zip_info.filename = path
        zip_info.date_time = (2022, 1, 1, 0, 0, 0)
        zip_info.external_attr = (stat.S_IFREG | permission) << 16
        self.zipfile.writestr(
            zip_info,
            contents,
            compress_type=ZIP_DEFLATED,
            compresslevel=9,
        )


class BaseAutomation:
    short_name: str = "base"
    description: str = "Abstract base automation."
    frequency: int
    working_dir: Optional[str]

    """
    The short name of the automation provider. This is used as an opaque token throughout
    the portal system.
    """

    def __init__(self, working_dir: Optional[str] = None):
        super().__init__()
        self.working_dir = working_dir

    def automate(self, full: bool = False) -> Tuple[bool, str]:
        raise NotImplementedError()

    def tmpl_write(self, filename: str, template: str, **kwargs: Any) -> None:
        """
        Write a Jinja2 template to the working directory for use by an automation module.

        :param filename: filename to write to
        :param template: Jinja2 template
        :param kwargs: variables for use with the template
        :return: None
        """
        if not self.working_dir:
            raise RuntimeError("No working directory specified.")
        tmpl = jinja2.Template(template)
        with open(
            os.path.join(self.working_dir, filename), "w", encoding="utf-8"
        ) as tfconf:
            tfconf.write(tmpl.render(**kwargs))

    def bin_write(
        self, filename: str, data: bytes, group_id: Optional[int] = None
    ) -> None:
        if not self.working_dir:
            raise RuntimeError("No working directory specified.")
        try:
            os.mkdir(os.path.join(self.working_dir, str(group_id)))
        except FileExistsError:
            pass
        with open(
            os.path.join(self.working_dir, str(group_id) if group_id else "", filename),
            "wb",
        ) as binfile:
            binfile.write(data)