From 2ac2e96c72672fcc20a6260eb458d0541d038529 Mon Sep 17 00:00:00 2001 From: renaud gaudin Date: Fri, 2 Dec 2022 15:33:17 +0000 Subject: [PATCH] Initial working version This is a working prototype that can be used to create working images from valid config files. Fixes: - Specify the card/project configuration format & workflow #4 - Implement Basic program structure #7 - Implement image file manipulation #10 Starts without completing: - Implement Progress Reporting #8: text UI (cli) is done and usable. Machine-readable JSON file is not but can be implemented via the logger and StepMachine. - Implement Downloader #9: a basic requests-based downloader is present. Needs more work to support transfer issues and being interupted TBC --- .github/workflows/qa.yml | 2 +- README.md | 170 ++++++------ requirements.txt | 6 + src/image_creator/__init__.py | 1 + src/image_creator/__main__.py | 3 + src/image_creator/constants.py | 91 ++++++ src/image_creator/creator.py | 40 +++ src/image_creator/entrypoint.py | 84 ++++++ src/image_creator/inputs.py | 210 ++++++++++++++ src/image_creator/logger.py | 185 +++++++++++++ src/image_creator/steps/__init__.py | 134 +++++++++ src/image_creator/steps/base.py | 101 +++++++ src/image_creator/steps/check_inputs.py | 195 +++++++++++++ src/image_creator/steps/contents.py | 354 ++++++++++++++++++++++++ src/image_creator/steps/image.py | 143 ++++++++++ src/image_creator/steps/oci_images.py | 37 +++ src/image_creator/steps/sizes.py | 24 ++ src/image_creator/utils/__init__.py | 0 src/image_creator/utils/download.py | 109 ++++++++ src/image_creator/utils/image.py | 298 ++++++++++++++++++++ src/image_creator/utils/misc.py | 94 +++++++ src/image_creator/utils/oci_images.py | 25 ++ src/image_creator/utils/requirements.py | 69 +++++ 23 files changed, 2296 insertions(+), 79 deletions(-) create mode 100644 requirements.txt create mode 100644 src/image_creator/__init__.py create mode 100644 src/image_creator/__main__.py create mode 100644 src/image_creator/constants.py create mode 100644 src/image_creator/creator.py create mode 100644 src/image_creator/entrypoint.py create mode 100644 src/image_creator/inputs.py create mode 100644 src/image_creator/logger.py create mode 100644 src/image_creator/steps/__init__.py create mode 100644 src/image_creator/steps/base.py create mode 100644 src/image_creator/steps/check_inputs.py create mode 100644 src/image_creator/steps/contents.py create mode 100644 src/image_creator/steps/image.py create mode 100644 src/image_creator/steps/oci_images.py create mode 100644 src/image_creator/steps/sizes.py create mode 100644 src/image_creator/utils/__init__.py create mode 100644 src/image_creator/utils/download.py create mode 100644 src/image_creator/utils/image.py create mode 100644 src/image_creator/utils/misc.py create mode 100644 src/image_creator/utils/oci_images.py create mode 100644 src/image_creator/utils/requirements.py diff --git a/.github/workflows/qa.yml b/.github/workflows/qa.yml index 3d35326..6da202a 100644 --- a/.github/workflows/qa.yml +++ b/.github/workflows/qa.yml @@ -15,7 +15,7 @@ jobs: - name: Setup Python uses: actions/setup-python@v2.3.3 with: - python-version: 3.10 + python-version: "3.10" architecture: x64 - name: Check black formatting run: | diff --git a/README.md b/README.md index 9d19a4e..601a77e 100644 --- a/README.md +++ b/README.md @@ -1,80 +1,94 @@ # image-creator -Hotspot image creator to build OLIP or Kiwix Offspot off [`base-image`](https://github.com/offspot/base-image). - -## Scope - -- Validate inputs -- Download base image -- Resize image to match contents -- Download contents into mounted `/data` -- Post-process downloaded contents -- Configure from inputs -- Re-generate SSH server keys -- *Pull* application images -- Prepares JSON config - -## Inputs - -- Target system (OLIP or Offspot) -- Image name -- Hostname -- domain name -- SSID -- WiFi AP password (if any) -- WiFi Country code -- WiFi channel -- Timezone -- SSH Public keys to add -- VPN configuration (tinc) -- Contents - -## App Containers - -- **OLIP** - - API - - Frontend - - Stats - - Controller -- **Offspot** - - Kiwix-serve - - WikiFundi (en/fr/es) - - Aflatoun (en/fr) - - Surfer -- IPFS daemon -- Captive portal - -## data partition - - -| /data subfolders | Usage | -|---|---| -| `offspot/zim` | Offspot Kiwix serve ZIM files| -| `offspot/wikifundi` | Offspot WikiFundi data | -| `offspot/files` | Offspot Surfer data | -| `offspot/xxx` | Offspot data for other apps | -| `olip` | OLIP data | - - -## JSON Configurator - -JSON config file at `/boot/config.json` is read and parsed on startup by the boot-time config script. -It looks for the following properties. Dotted ones means nested. - -Behavior is to adjust configuration only if the property is present. Script will remove property from JSON once applied. - -Configurator is also responsible for resizing `/data` partition to device size on first boot but this is not configurable via JSON. - -| Property| Type | Usage | -|---|---|---| -| `hostname` | `string` | Pi host name | -| `domain` | `string` | FQDN to answer to on DNS | -| `wifi.ssid` | `string` | WiFi SSID | -| `wifi.password` | `string` | WiFi password (clear). If `null`, auth not required | -| `wifi.country-code` | `string` | ISO-639-2 Country code for WiFI | -| `wifi.channel` | `int` | 1-11 channel for WiFi | -| `timezone` | `string` | Timezone to configure date with | -| `ssh-keys` | `string[]` | List of public keys to add to user | -| `tinc-vpn` | `string` | tinc-VPN configuration | -| `env.all` | `string[]` | List of `KEY=VALUE` environment variables to pass to **all applications** | -| `env.xxx` | `string[]` | List of `KEY=VALUE` environment variables to pass **containers matching _xxx_** | \ No newline at end of file +RaspberryPi image creator to build OLIP or Kiwix Hotspot off [`base-image`](https://github.com/offspot/base-image). + +[![CodeFactor](https://www.codefactor.io/repository/github/offspot/image-creator/badge)](https://www.codefactor.io/repository/github/offspot/image-creator) +[![Build Status](https://github.com/offspot/image-creator/actions/workflows/build.yml/badge.svg?branch=main)](https://github.com/offspot/image-creator/actions/workflows/build.yml?query=branch%3Amain) +[![License: GPL v3](https://img.shields.io/badge/License-GPLv3-blue.svg)](https://www.gnu.org/licenses/gpl-3.0) + + +## Usage + +`image-creator` is to be **ran as `root`**. + +``` +❯ image-creator --help +usage: image-creator [-h] [--build-dir BUILD_DIR] [-C] [-K] [-X] [-T CONCURRENCY] [-D] [-V] CONFIG_SRC OUTPUT + +create an Offspot Image from a config + +positional arguments: + CONFIG_SRC Offspot Config YAML file path or URL + OUTPUT Where to write image to + +options: + -h, --help show this help message and exit + --build-dir BUILD_DIR + Directory to store temporary files in, like files that needs to be extracted. Defaults to some place within /tmp + -C, --check Only check inputs, URLs and sizes. Don't download/create image. + -K, --keep [DEBUG] Don't remove output image if creation failed + -X, --overwrite Don't fail on existing output image: remove instead + -T CONCURRENCY, --concurrency CONCURRENCY + Nb. of threads to start for parallel downloads (at most one per file). `0` (default) for auto-selection based on CPUs. + `1` to disable concurrency. + -D, --debug + -V, --version show program's version number and exit +``` + + +## Configuration + +Image configuration is done through a YAML file which must match the following format. Only `base` is required. + + + +| Member | Kind | Function | +|------------------|----------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `base` | `string` | Version ([official releases](https://drive.offspot.it/base/)) or URL to a base-image file. Accepts `file://` URLs. Accepts lzma encoded images using `.xz` suffix | +| `output.size` | `string`/`int` | Requested size of output image. Accepts `auto` for an power-of-2 sized that can fit the content (⚠️ TBI) | +| `oci_images` | `string[]` | List of OCI Image names. More specific the better. Prefer ghcr.io if possible | +| `files` | `file[]` | List of files to include on the data partition. See below. One of `url` or `content` must be present | +| `files[].url` | `string` | URL to download file from | +| `files[].to` | `string` | [required] Path to store file at. Must be a descendent of `/data` | +| `files[].content`| `string` | Text content of the file to write. Replaces `url` if present | +| `files[].via` | `string` | For `url`-based files, transformation to apply on downloaded file: `direct` (default): simple download, `bztar`, `gztar`, `tar`, `xztar`, `zip` to expand archives | +| `files[].size` | `string`/`int` | **Only for `untar`/`unzip`** should file be compressed. Specify expanded size. Assumes File-size (uncompressed) if not specified. ⚠️ Fails if lower than file size | +| `write_config` | `bool` | Whether to write this file to `/data/conf/image.yaml` | +| `offspot` | `dict` | [runtime-config](https://github.com/offspot/runtime-config) configuration. Will be parsed and dumped to `/boot/offspot.yaml` | + +### Sample + +```yaml +--- +base: 1.0.0 +output: + size: 8G +oci_images: +- ghcr.io/offspot/kiwix-serve:dev +files: +- url: http://download.kiwix.org/zim/wikipedia_fr_test.zim + to: /data/contents/zims/wikipedia_fr_test.zim + via: direct +- to: /data/conf/message.txt + content: | + hello world +wite_config: true +offspot: + timezone: Africa/Bamako + ap: + ssid: Kiwix Offspot + as-gateway: true + domain: demo + tld: offspot + containers: + services: + kiwix: + container_name: kiwix + image: ghcr.io/offspot/kiwix-serve:dev + command: /bin/sh -c "kiwix-serve /data/*.zim" + volumes: + - "/data/content/zims:/data:ro" + ports: + - "80:80" + +``` diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..37f6cdd --- /dev/null +++ b/requirements.txt @@ -0,0 +1,6 @@ +requests==2.28.1 +PyYAML==6.0 +cli-ui==0.17.2 +humanfriendly==10.0 +progressbar2==4.2.0 +docker_export==0.4 diff --git a/src/image_creator/__init__.py b/src/image_creator/__init__.py new file mode 100644 index 0000000..e3ac18c --- /dev/null +++ b/src/image_creator/__init__.py @@ -0,0 +1 @@ +__version__ = "1.0.0.dev0" diff --git a/src/image_creator/__main__.py b/src/image_creator/__main__.py new file mode 100644 index 0000000..591ab8e --- /dev/null +++ b/src/image_creator/__main__.py @@ -0,0 +1,3 @@ +from image_creator.entrypoint import main + +main() diff --git a/src/image_creator/constants.py b/src/image_creator/constants.py new file mode 100644 index 0000000..2321e09 --- /dev/null +++ b/src/image_creator/constants.py @@ -0,0 +1,91 @@ +import logging +import pathlib +import re +import sys +import tempfile +import urllib.parse +from dataclasses import dataclass +from typing import Union + +from image_creator import __version__ as vers +from image_creator.logger import Logger + +# where will data partition be monted on final device. +# used as reference for destinations in config file and in the UI +DATA_PART_PATH = pathlib.Path("/data") +# version of the python interpreter +pyvers = ".".join([str(p) for p in sys.version_info[:3]]) +banner: str = rf""" + _ _ + (_)_ __ ___ __ _ __ _ ___ ___ _ __ ___ __ _| |_ ___ _ __ + | | '_ ` _ \ / _` |/ _` |/ _ \_____ / __| '__/ _ \/ _` | __/ _ \| '__| + | | | | | | | (_| | (_| | __/_____| (__| | | __/ (_| | || (_) | | + |_|_| |_| |_|\__,_|\__, |\___| \___|_| \___|\__,_|\__\___/|_| + |___/ v{vers}|py{pyvers} + +""" + + +@dataclass(kw_only=True) +class Options: + """Command-line options""" + + CONFIG_SRC: str + OUTPUT: str + BUILD_DIR: str + + check_only: bool + debug: bool + + config_path: pathlib.Path = None + output_path: pathlib.Path = None + build_dir: pathlib.Path = None + + keep_failed: bool + overwrite: bool + concurrency: int + + config_url: urllib.parse.ParseResult = None + logger: Logger = Logger() + + def __post_init__(self): + if re.match(r"^https?://", self.CONFIG_SRC): + self.config_url = urllib.parse.urlparse(self.CONFIG_SRC) + else: + self.config_path = pathlib.Path(self.CONFIG_SRC).expanduser().resolve() + + if self.debug: + self.logger.setLevel(logging.DEBUG) + + if not self.check_only: + self.output_path = pathlib.Path(self.OUTPUT).expanduser().resolve() + + if not self.BUILD_DIR: + # holds reference to tempdir until Options is released + # and will thus automatically remove actual folder + self.__build_dir = tempfile.TemporaryDirectory( + prefix="image-creator_build-dir", ignore_cleanup_errors=True + ) + self.build_dir = ( + pathlib.Path(self.BUILD_DIR or self.__build_dir.name).expanduser().resolve() + ) + + @property + def version(self): + return vers + + @property + def config_src(self) -> Union[pathlib.Path, urllib.parse.ParseResult]: + return self.config_url or self.config_path + + +class _Global: + options = None + + @property + def logger(self): + return Global.options.logger if Global.options else Options.logger + + +Global = _Global() +logger = Global.logger diff --git a/src/image_creator/creator.py b/src/image_creator/creator.py new file mode 100644 index 0000000..fe1f685 --- /dev/null +++ b/src/image_creator/creator.py @@ -0,0 +1,40 @@ +import atexit + +from image_creator.constants import Global, Options, banner, logger +from image_creator.logger import Status +from image_creator.steps import StepMachine +from image_creator.utils.misc import rmtree + + +class ImageCreator: + def __init__(self, **kwargs): + Global.options = Options(**kwargs) + # make sure we clean things up before exiting + atexit.register(self.halt) + + def run(self): + if Global.options.check_only: + StepMachine.halt_after("sizes:ComputeSizes") + + logger.message(banner) + + self.machine = StepMachine(options=Global.options) + for step in self.machine: + logger.start_step(step.name) + res = step.execute() + logger.end_step() + if res != 0: + logger.error(f"Step “{repr(step)}” returned {res}") + return res + + def halt(self): + logger.message("Cleaning-up…", end=" ", timed=True) + self.machine.halt() + if Global.options.build_dir: + try: + rmtree(Global.options.build_dir) + except Exception: + logger.add_dot(Status.NOK) + else: + logger.add_dot(Status.OK) + logger.message() diff --git a/src/image_creator/entrypoint.py b/src/image_creator/entrypoint.py new file mode 100644 index 0000000..195c74c --- /dev/null +++ b/src/image_creator/entrypoint.py @@ -0,0 +1,84 @@ +import argparse +import sys +import tempfile + +from image_creator import __version__ +from image_creator.constants import logger +from image_creator.creator import ImageCreator + + +def main(): + parser = argparse.ArgumentParser( + prog="image-creator", + description="create an Offspot Image from a config", + ) + + parser.add_argument( + "--build-dir", + dest="BUILD_DIR", + help="Directory to store temporary files in, " + "like files that needs to be extracted. " + f"Defaults to some place within {tempfile.gettempdir()}", + ) + parser.add_argument( + "-C", + "--check", + action="store_true", + dest="check_only", + help="Only check inputs, URLs and sizes. Don't download/create image.", + ) + + parser.add_argument( + "-K", + "--keep", + action="store_true", + dest="keep_failed", + default=False, + help="[DEBUG] Don't remove output image if creation failed", + ) + parser.add_argument( + "-X", + "--overwrite", + action="store_true", + default=False, + dest="overwrite", + help="Don't fail on existing output image: remove instead", + ) + parser.add_argument( + "-T", + "--concurrency", + type=int, + default=0, + dest="concurrency", + help="Nb. of threads to start for parallel downloads (at most one per file). " + "`0` (default) for auto-selection based on CPUs. `1` to disable concurrency.", + ) + parser.add_argument("-D", "--debug", action="store_true", dest="debug") + parser.add_argument("-V", "--version", action="version", version=__version__) + + parser.add_argument(help="Offspot Config YAML file path or URL", dest="CONFIG_SRC") + parser.add_argument( + dest="OUTPUT", + help="Where to write image to", + ) + + kwargs = dict(parser.parse_args()._get_kwargs()) + + try: + app = ImageCreator(**kwargs) + sys.exit(app.run()) + except Exception as exc: + if kwargs.get("debug"): + logger.exception(exc) + logger.critical(str(exc)) + try: + app.halt() + except Exception as exc: + logger.debug(f"Errors cleaning-up: {exc}") + sys.exit(1) + finally: + logger.terminate() + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/image_creator/inputs.py b/src/image_creator/inputs.py new file mode 100644 index 0000000..f0696c0 --- /dev/null +++ b/src/image_creator/inputs.py @@ -0,0 +1,210 @@ +import pathlib +import re +import shutil +import urllib.parse +from typing import Any, Dict, List, Optional, Tuple, Union + +try: + from yaml import CLoader as Loader + from yaml import load as yaml_load +except ImportError: + # we don't NEED cython ext but it's faster so use it if avail. + from yaml import Loader, load as yaml_load + +from image_creator.constants import DATA_PART_PATH +from image_creator.utils.download import get_online_rsc_size +from image_creator.utils.misc import get_filesize +from image_creator.utils.oci_images import Image as OCIImage + + +def has_val(data, key: str) -> bool: + value = data.get(key, "") + if not value: + return False + if not isinstance(value, str): + return False + return bool(value) + + +class File: + """In-Config reference to a file to write to the data partition + + Created from files entries in config: + - to: str + mandatory destination to save file into. Must be inside /data + - size: optional[int] + size of (expanded) content. If specified, must be >= source file + - via: optional[str] + method to process source file (not for content). Values in File.unpack_formats + - url: optional[str] + URL to download file from + - content: optional[str] + plain text content to write to destination + + one of content or url must be supplied. content has priority""" + + unpack_formats = [f[0] for f in shutil.get_unpack_formats()] + + def __init__(self, payload: Dict[str, Union[str, int]]): + self.url = None + self.content = payload.get("content") + + if not self.content: + try: + self.url = urllib.parse.urlparse(payload.get("url")) + except Exception: + raise ValueError(f"URL “{payload.get('url')}” is incorrect") + + self.to = pathlib.Path(payload["to"]).resolve() + if not self.to.is_relative_to(DATA_PART_PATH): + raise ValueError(f"{self.to} not a descendent of {DATA_PART_PATH}") + + self.via = payload.get("via", "direct") + if self.via not in ("direct", "unzip", "untar"): + raise NotImplementedError(f"Unsupported handler `{self.via}`") + + # initialized has unknown + self.size = payload.get("size", -1) + + def fetch_size(self, force: Optional[bool] = False) -> int: + """retrieve size of source, making sure it's reachable""" + if not force and self.size >= 0: + return self.size + self.size = ( + get_filesize(self.getpath()) + if self.is_local + else get_online_rsc_size(self.geturl()) + ) + return self.size + + def geturl(self) -> str: + """URL as string""" + try: + return self.url.geturl() + except Exception: + return None + + def getpath(self) -> pathlib.Path: + """URL as a local path""" + return pathlib.Path(self.url.path).expanduser().resolve() + + @property + def is_direct(self): + return self.via == "direct" + + @property + def is_plain(self) -> bool: + """whether a plain text content to be written""" + return self.content is not None + + @property + def is_local(self) -> bool: + """whether referencing a local file""" + return not self.is_plain and self.url and self.url.scheme == "file" + + @property + def is_remote(self) -> bool: + """whether referencing a remote file""" + return self.content is None and self.url and self.url.scheme != "file" + + def mounted_to(self, mount_point: pathlib.Path): + """destination (to) path inside mount-point""" + return mount_point.joinpath(self.to.relative_to(DATA_PART_PATH)) + + def __repr__(self) -> str: + msg = f"File(to={self.to}, via={self.via}" + if self.url: + msg += f", url={self.geturl()}" + if self.content: + msg += f", content={self.content.splitlines()[0][:10]}" + msg += f", size={self.size})" + return msg + + def __str__(self) -> str: + return repr(self) + + +class Config(dict): + """Parsed Image YAML Configuration""" + + @classmethod + def read_from(cls, text: str): + """Instanciate from yaml text""" + return cls(**yaml_load(text, Loader=Loader)) + + def init(self): + """Prepare Config from yaml-parsed dict""" + self.errors: List[Tuple[str, str]] = [] + + self.base: File = self._get_base() + self.all_files: List[File] = [ + File(payload) for payload in self.get("files", []) + ] + self.oci_images: List[OCIImage] = [ + OCIImage.parse(str(name)) for name in self.get("oci_images", []) + ] + return self.validate() + + @property + def is_valid(self) -> bool: + return not self.errors + + def _get_base(self) -> File: + """Infer url from flexible `base` and return a File""" + url = self.get("base") + match = re.match(r"^(?P\d\.\d\.\d)(?P[a-z0-9\-\.\_]*)", url) + if match: + version = "".join(match.groups()) + url = f"https://drive.offspot.it/base/base-image-{version}.img.xz" + return File({"url": url, "to": DATA_PART_PATH / "-"}) + + def dig(self, path, default=None) -> Any: + """get a value using it's dotted tree path""" + data = self + parts = path.split(".") + for index in range(0, len(parts) - 1): + data = data.get(parts[index], {}) + return data.get(parts[-1], default) + + @property + def offspot_config(self) -> Dict: + """parsed `offspot` subtree representing runtime-config file""" + return self.get("offspot") + + @property + def remote_files(self) -> List[File]: + return [file for file in self.all_files if file.is_remote] + + @property + def non_remote_files(self) -> List[File]: + return [file for file in self.all_files if file.is_plain or file.is_local] + + def validate(self) -> bool: + """whether Config can be run or not + + Feedback for user in self.errors""" + + # check for required props (only base ATM) + for key in ("base",): + if not self.get(key): + self.errors.append(key, f"missing `{key}`") + + # check that files are OK + files = self.get("files", []) + if not isinstance(files, list): + self.errors.append(("files", "not a list")) + + for file in files: + if not isinstance(file, dict): + self.errors.append("files", "not a dict") + if not has_val(file, "to"): + self.errors.append("files.to", "missing or invalid") + if not has_val(file, "url") and not has_val(file, "content"): + self.errors.append("files", "`url` or `content` must be set") + + # make sure no two-files have the same destination + all_tos = [file.get("to") for file in files] + if len(all_tos) != len(set(all_tos)): + self.errors.append(("files", "using same `to:` target several times")) + + return self.is_valid diff --git a/src/image_creator/logger.py b/src/image_creator/logger.py new file mode 100644 index 0000000..cc37cd4 --- /dev/null +++ b/src/image_creator/logger.py @@ -0,0 +1,185 @@ +import enum +import logging +import pathlib +import traceback +from typing import Optional + +import cli_ui as ui + +Status = enum.Enum("Status", ["OK", "NOK", "NEUTRAL"]) +Colors = {Status.NEUTRAL: ui.reset, Status.OK: ui.green, Status.NOK: ui.red} +ui.warn = ui.UnicodeSequence(ui.brown, "⚠️", "/!\\") + + +class Logger: + """Custom cli_ui-based logger providing ~unified UI for steps and tasks + + Operations are either Steps or tasks within a Step + - Steps are collections of tasks + - UI displays Steps and Tasks differently using symbols and indentation + - Most tasks are visually represented based on ~state (not recorded) + - running: started but not ended) + - suceeded: ended successfuly ; with an optional confirmation text + - failed: ended unsuccessfuly ; providing reason + + Most of this logger's job is to abstract this visual organization behind a + hierachical API: + - start_step() + - start_task() + - succeed_task() + - end_step() + """ + + def __init__( + self, + level: Optional[int] = logging.INFO, + progress_to: Optional[pathlib.Path] = None, + ): + self.verbose = level + self.progress_to = progress_to + + self.setLevel(level) + + self.currently = None + + def setLevel(self, level: int): + """reset logger's verbose config based on level""" + ui.setup( + verbose=level <= logging.DEBUG, + quiet=level >= logging.WARNING, + color="auto", + title="image-creator", + timestamp=False, + ) + + def message(self, *tokens, end: str = "\n", timed: bool = False): + """Flexible message printing + + - end: controls carriage-return + - timed: control wehther to prefix with time""" + self.clear() + if timed: + ui.CONFIG["timestamp"] = True + ui.message(" " * self.indent_level, *tokens, end=end) + if timed: + ui.CONFIG["timestamp"] = False + + def debug(self, text: str): + self.clear() + ui.debug(ui.indent(text, num=self.indent_level)) + + def info(self, text: str, end: str = "\n"): + self.clear() + ui.info(ui.indent(text, num=self.indent_level), end=end) + + def warning(self, text: str): + self.clear() + ui.message(ui.brown, ui.indent(text, num=self.indent_level)) + + def error(self, text: str): + self.clear() + ui.message(ui.bold, ui.red, ui.indent(text, num=self.indent_level)) + + def exception(self, exc: Exception): + ui.message(ui.red, "".join(traceback.format_exception(exc))) + + def critical(self, text: str): + self.clear() + ui.error(text) + + def fatal(self, text: str): + self.critical(text) + + @property + def with_progress(self) -> bool: + """wether configured to write progress to an external machine-readable file""" + return self.progress_to is not None + + @property + def indent_level(self): + """standard indentation level based on current ~position""" + match self.currently: + case "step": + return 3 + case "task": + return 6 + case _: + return 0 + + def mark_as(self, what: str): + """set new ~position of the logger: step or task""" + self.currently = what + + def clear(self): + """clear in-task or in-step same-line hanging to prevent writing to previous""" + if self.currently in ("task",): + ui.info("") + + def start_step(self, step: str): + """Start a new Step, Step has no status and will be on a single line""" + self.clear() + self.mark_as("step") + ui.CONFIG["timestamp"] = True + ui.info_1(step) + ui.CONFIG["timestamp"] = False + + def end_step(self): + self.clear() + self.mark_as(None) + + def start_task(self, task: str): + """Start new task. Task is expectd to end""" + self.clear() + self.mark_as("task") + ui.CONFIG["timestamp"] = True + ui.message(" ", ui.bold, ui.blue, "=>", ui.reset, task, end=" ") + ui.CONFIG["timestamp"] = False + + def end_task(self, success: Optional[bool] = None, message: Optional[str] = None): + """End current task with custom success symbol and message""" + tokens = [] if success is None else [ui.check if success else ui.cross] + if message: + tokens += [ui.brown, message] + ui.message(*tokens) + self.mark_as(None) + + def succeed_task(self, message: Optional[str] = None): + """End current task as successful with optional message""" + self.end_task(success=True, message=message) + + def fail_task(self, message: Optional[str] = None): + """End current task as unsuccessful with optional message""" + self.end_task(success=False, message=message) + + def add_task(self, name: str, message: str = None): + """Single-call task with no status information""" + self.start_task(name) + if message: + ui.message(*[ui.brown, message]) + else: + ui.message() + self.mark_as(None) + + def complete_download( + self, + name: str, + size: Optional[str] = None, + extra: Optional[str] = None, + failed: Optional[bool] = False, + ): + """record completed download, inside a task, potentially following progress""" + tokens = [" ", ui.warn if failed else ui.check, name] + if size: + tokens += [ui.brown, str(size)] + if extra: + tokens += [ui.reset, extra] + self.message(*tokens, timed=True) + + def add_dot(self, status: Status = Status.NEUTRAL): + """pytest-like colored dots indicating hidden operations status + + Must be cleared-out manually with a newline (ui.message())""" + ui.message(Colors.get(status, Colors[Status.NEUTRAL]), ".", end="") + + def terminate(self): + self.clear() diff --git a/src/image_creator/steps/__init__.py b/src/image_creator/steps/__init__.py new file mode 100644 index 0000000..a5eb690 --- /dev/null +++ b/src/image_creator/steps/__init__.py @@ -0,0 +1,134 @@ +import sys +from importlib import import_module +from typing import Any, Dict + +from image_creator.constants import logger +from image_creator.logger import Status + + +class StepMachine: + """Ordered list of Steps""" + + steps = [ + ".:VirtualInitStep", + "check_inputs:CheckRequirements", + "check_inputs:CheckInputs", + "check_inputs:CheckURLs", + "sizes:ComputeSizes", + # check-only stops here + "base:DownloadImage", + "image:ResizingImage", + "image:MountingDataPart", + "oci_images:DownloadingOCIImages", + "contents:ProcessingLocalContent", + "contents:DownloadingContent", + "image:UnmountingDataPart", + "image:MountingBootPart", + "check_inputs:WritingOffspotConfig", + "image:UnmountingBootPart", + "image:DetachingImage", + ] + + def __init__(self, **kwargs): + self.payload = dict(**kwargs) + self._current = 0 + self.step = self._get_step(self._current) + + @classmethod + def halt_after(cls, step: str): + """reduce StepMachine to end with that step""" + cls.steps = cls.steps[: cls.steps.index(step) + 1] + + def _get_step(self, index: int): + modname, target = self.steps[index].split(":", 1) + module = ( + sys.modules[__name__] + if modname == "." + else import_module(f"image_creator.steps.{modname}") + ) + return getattr(module, target).__call__(self) + + @property + def step_num(self): + """current step number (1-indexed)""" + return self._current + 1 + + def __iter__(self): + return self + + def __next__(self): + try: + new_index = self._current + 1 + self.steps[new_index] + except IndexError: + raise StopIteration() + + try: + self.step = self._get_step(new_index) + except Exception as exc: + logger.error(f"failed to init step {self.steps[new_index]}: {exc}") + logger.exception(exc) + raise StopIteration + self._current = new_index + return self.step + + def halt(self): + """request cleanup of ran-steps + + Steps being dependent on the previous ones, some resources such as + loop-device or mount-points are passed along several steps + this calls individual cleanup() methods in reversed order""" + + # calling .cleanup() on each step from last called one, in reverse order + for index in range(self._current, 0, -1): + step = self._get_step(index) + try: + step.cleanup(payload=self.payload) + except Exception: + logger.add_dot(status=Status.NOK) + else: + logger.add_dot(status=Status.OK) + if ( + not self.payload["options"].keep_failed + and self.payload["options"].output_path + and self.payload["options"].output_path.exists() + ): + try: + self.payload["options"].output_path.unlink(missing_ok=True) + except Exception: + logger.add_dot(status=Status.NOK) + else: + logger.add_dot(status=Status.OK) + + +class Step: + """StepInterface""" + + def __init__(self, machine: StepMachine): + self.machine = machine + + # name of step to be overriden + @property + def name(self): + return repr(self) + + def __repr__(self): + return self.__class__.__name__ + + def __str__(self): + return self.name + + def execute(self, *args, **kwargs) -> int: + return self.run(self.machine.payload, *args, **kwargs) + + def run(self, payload: Dict[str, Any], *args, **kwargs) -> int: + """actual step implementation. 0 on success""" + raise NotImplementedError() + + def cleanup(self, payload): + """clean resources reserved in run()""" + ... + + +class VirtualInitStep(Step): + ... diff --git a/src/image_creator/steps/base.py b/src/image_creator/steps/base.py new file mode 100644 index 0000000..4b842f6 --- /dev/null +++ b/src/image_creator/steps/base.py @@ -0,0 +1,101 @@ +import shutil + +from image_creator.constants import logger +from image_creator.steps import Step +from image_creator.utils.download import download_file +from image_creator.utils.misc import extract_xz_image, format_size, get_filesize + + +class DownloadImage(Step): + name = "Fetching base image" + + def run(self, payload, *args, **kwargs) -> int: + """whether system requirements are satisfied""" + base_file = payload["config"].base + + # we'll have to download it + if not base_file.is_local: + rem_path = base_file.getpath() + + # we'll have to download it to build-dir then extract + if rem_path.suffix == ".xz": + xz_fpath = payload["options"].build_dir.joinpath(rem_path.name) + logger.start_task(f"Downloading {base_file} into {xz_fpath}…") + try: + download_file(base_file.geturl(), xz_fpath) + except Exception as exc: + logger.fail_task(str(exc)) + return 1 + else: + logger.succeed_task(format_size(get_filesize(xz_fpath))) + + logger.start_task( + f"Extracting {xz_fpath} into {payload['options'].output_path}…" + ) + try: + extract_xz_image(xz_fpath, payload["options"].output_path) + except Exception as exc: + logger.fail_task(str(exc)) + return 1 + else: + logger.succeed_task( + format_size(get_filesize(payload["options"].output_path)) + ) + + logger.start_task(f"Removing {xz_fpath}…") + try: + xz_fpath.unlink(missing_ok=True) + except Exception as exc: + logger.fail_task(str(exc)) + else: + logger.succeed_task() + + # download straight to final destination + else: + logger.start_task( + f"Downloading {base_file} into {payload['options'].output_path}…" + ) + try: + download_file(base_file.geturl(), payload["options"].output_path) + except Exception as exc: + logger.fail_task(str(exc)) + return 1 + else: + logger.succeed_task( + format_size(get_filesize(payload["options"].output_path)) + ) + # no need to download + else: + # we'll extract it to destination directly + if base_file.getpath().suffix == ".xz": + logger.start_task( + f"Extracting {base_file.getpath()} " + f"into {payload['options'].output_path}…" + ) + try: + extract_xz_image( + base_file.getpath(), payload["options"].output_path + ) + except Exception as exc: + logger.fail_task(str(exc)) + return 1 + else: + logger.succeed_task( + format_size(get_filesize(payload["options"].output_path)) + ) + # we'll simply copy it to destination + else: + logger.start_task( + f"Copying {base_file.getpath()} " + f"into {payload['options'].output_path}…" + ) + try: + shutil.copy2(base_file.getpath(), payload["options"].output_path) + except Exception as exc: + logger.fail_task(str(exc)) + return 1 + else: + logger.succeed_task( + format_size(get_filesize(payload["options"].output_path)) + ) + return 0 diff --git a/src/image_creator/steps/check_inputs.py b/src/image_creator/steps/check_inputs.py new file mode 100644 index 0000000..929bd11 --- /dev/null +++ b/src/image_creator/steps/check_inputs.py @@ -0,0 +1,195 @@ +import pathlib + +try: + from yaml import CDumper as Dumper + from yaml import dump as yaml_dump +except ImportError: + # we don't NEED cython ext but it's faster so use it if avail. + from yaml import Dumper, dump as yaml_dump + +from image_creator.constants import logger +from image_creator.inputs import Config +from image_creator.steps import Step +from image_creator.utils import requirements +from image_creator.utils.download import read_text_from +from image_creator.utils.misc import format_size +from image_creator.utils.oci_images import image_exists + + +class CheckRequirements(Step): + name = "Checking system requirements…" + + def run(self, payload, *args, **kwargs) -> int: + """whether system requirements are satisfied""" + + all_good = True + + logger.start_task("Checking uid…") + if not requirements.is_root(): + logger.fail_task("you must be root") + all_good &= False + else: + logger.succeed_task() + + logger.start_task("Checking binary dependencies") + has_all, missing_bins = requirements.has_all_binaries() + if not has_all: + all_good &= False + logger.fail_task(f"Missing binaries: {', '.join(missing_bins)}") + else: + logger.succeed_task() + + logger.start_task("Checking loop-device capability") + if not requirements.has_loop_device(): + all_good &= False + logger.fail_task() + else: + logger.succeed_task() + + logger.start_task("Checking ext4 support") + if not requirements.has_ext4_support(): + all_good &= False + logger.fail_task() + else: + logger.succeed_task() + + if not all_good: + logger.warning(requirements.help_text) + + return 0 if all_good else 2 + + +class CheckInputs(Step): + name = "Checking config inputs…" + + def run(self, payload, *args, **kwargs) -> int: + logger.start_task(f"Reading config from {payload['options'].config_src}") + + try: + if isinstance(payload["options"].config_src, pathlib.Path): + text = payload["options"].config_src.read_text() + else: + text = read_text_from(payload["options"].config_src) + except Exception as exc: + logger.fail_task(str(exc)) + raise exc + else: + logger.succeed_task() + + logger.start_task("Parsing config data…") + try: + payload["config"] = Config.read_from(text) + except Exception as exc: + logger.fail_task() + raise exc + else: + logger.succeed_task() + + logger.start_task("Checking parameters…") + try: + if not payload["config"].init(): + logger.fail_task("Config file is not valid") + logger.warning( + "\n".join( + [ + f"- [{key}] {error}" + for key, error in payload["config"].errors + ] + ) + ) + return 3 + else: + logger.succeed_task() + except Exception as exc: + logger.fail_task(f"Config contains invalid values: {exc}") + return 3 + + logger.start_task("Making sure base and output are different…") + if ( + payload["config"].base.is_local + and payload["config"].base.getpath() == payload["options"].output_path + ): + logger.fail_task("base and output image are the same") + return 3 + else: + logger.succeed_task() + + if payload["options"].output_path.exists() and payload["options"].overwrite: + logger.start_task("Removing target path…") + try: + payload["options"].output_path.unlink() + except Exception as exc: + logger.fail_task(str(exc)) + else: + logger.succeed_task() + else: + logger.start_task("Checking target path…") + if payload["options"].output_path.exists(): + logger.fail_task(f"{payload['options'].output_path} exists.") + return 3 + else: + logger.succeed_task() + + logger.start_task("Testing target location…") + try: + payload["options"].output_path.touch() + payload["options"].output_path.unlink() + except Exception as exc: + logger.fail_task(str(exc)) + else: + logger.succeed_task() + + return 0 + + +class CheckURLs(Step): + name = "Checking all Sources…" + + def run(self, payload, *args, **kwargs) -> int: + all_valid = True + + for file in [payload["config"].base] + payload["config"].all_files: + if file.is_plain: + continue + logger.start_task(f"Checking {file.geturl()}…") + size = file.fetch_size() + if size >= 0: + logger.succeed_task(format_size(size)) + elif size == -1: + logger.succeed_task("size unknown") + else: + logger.fail_task() + all_valid &= False + + for image in payload["config"].oci_images: + logger.start_task(f"Checking OCI Image {image}…") + if not image_exists(image): + logger.fail_task() + all_valid &= False + else: + logger.succeed_task() + + return 0 if all_valid else 4 + + +class WritingOffspotConfig(Step): + name = "Writing Offspot Config…" + + def run(self, payload, *args, **kwargs) -> int: + if not payload["config"].offspot_config: + logger.add_task("No Offspot config passed") + return 0 + + offspot_fpath = payload["image"].p1_mounted_on.joinpath("offspot.yaml") + logger.start_task(f"Saving Offspot config to {offspot_fpath}…") + try: + offspot_fpath.write_text( + yaml_dump(payload["config"].offspot_config, Dumper=Dumper) + ) + except Exception as exc: + logger.fail_task(str(exc)) + return 1 + else: + logger.succeed_task() + + return 0 diff --git a/src/image_creator/steps/contents.py b/src/image_creator/steps/contents.py new file mode 100644 index 0000000..99505f4 --- /dev/null +++ b/src/image_creator/steps/contents.py @@ -0,0 +1,354 @@ +import pathlib +import shutil +import tempfile +from collections import OrderedDict as od +from concurrent.futures import CancelledError, Future, ThreadPoolExecutor +from dataclasses import dataclass +from typing import Any, Callable, Optional + +import progressbar + +from image_creator.constants import logger +from image_creator.inputs import File +from image_creator.steps import Step +from image_creator.utils.download import download_file +from image_creator.utils.misc import ( + ensure_dir, + expand_file, + format_size, + get_filesize, + get_size_of, +) + + +@dataclass +class MultiDownloadProgress: + """Holds progress for multi-downloads""" + + nb_total: int = 0 + nb_completed: int = 0 + bytes_total: int = 0 + bytes_received: int = 0 + + def on_data(self, bytes_received: int): + self.bytes_received += bytes_received + + @property + def nb_remaining(self) -> int: + return self.nb_total - self.nb_completed + + @property + def bytes_remaining(self) -> int: + return self.bytes_total - self.bytes_received + + def __str__(self) -> str: + if not self.nb_total: + return repr(self) + + if not self.nb_remaining: + return f"{self.nb_completed} items completed" + + return ( + f"{self.nb_remaining} items accounting " + f"{format_size(self.bytes_remaining)} remaining" + ) + + +class MultiDownloadProgressBar: + """Custom progress bar tailored for MultiDownloadProgress + + Displays as: + + [Elapsed Time: 0:02:45] 1 MiB of 20 MiB downloaded|##### |1 KiB/s (Time: 0:00:45) + """ + + def __init__(self, dl_progress: MultiDownloadProgress): + widgets = [ + "[", + progressbar.Timer(), + "] ", + progressbar.DataSize(), + f" of {format_size(dl_progress.bytes_total)} downloaded", + progressbar.Bar(), + progressbar.AdaptiveTransferSpeed(), + " (", + progressbar.ETA(), + ")", + ] + self.bar = progressbar.ProgressBar( + max_value=dl_progress.bytes_total, widgets=widgets + ) + self.dl_progress = dl_progress + + def update(self): + self.bar.update(self.dl_progress.bytes_received) + + def finish(self): + self.bar.finish() + + +def download_file_worker( + file: File, + mount_point: pathlib.Path, + temp_dir: pathlib.Path, + on_data: Optional[Callable] = None, +): + """Downloads a File into its destination, unpacking if required""" + block_size = 2**20 # 1MiB + + dest_path = file.mounted_to(mount_point) + ensure_dir(dest_path.parent) + + if file.is_direct: + download_file( + file.geturl(), + dest_path, + block_size=block_size, + on_data=on_data, + ) + else: + temp_path = pathlib.Path( + tempfile.NamedTemporaryFile(dir=temp_dir, suffix=dest_path.name).name + ) + try: + download_file( + file.geturl(), + temp_path, + block_size=block_size, + on_data=on_data, + ) + except Exception as exc: + temp_path.unlink(missing_ok=True) + raise exc + try: + expand_file(temp_path, dest_path, file.via) + except Exception as exc: + raise exc + finally: + temp_path.unlink(missing_ok=True) + + +class FilesMultiDownloader: + """Downloads the File objects supplied using a ThreadPoolExecutor""" + + def __init__( + self, + files, + mount_point: pathlib.Path, + temp_dir: Optional[pathlib.Path] = None, + concurrency: Optional[int] = None, + callback: Optional[Callable] = None, + on_data: Optional[Callable] = None, + ): + self.files = files + self.mount_point = mount_point + self.temp_dir = temp_dir + + self.callback = callback + self.on_data = on_data + self.is_running = False + self.futures, self.cancelled, self.succeeded, self.failed = ( + od(), + od(), + od(), + od(), + ) + self.executor = ThreadPoolExecutor(max_workers=concurrency or None) + + def notify_completion(self, future: Future): + """called once Future id done. Moves it to approp. list then callback()""" + if not future.done(): + return + if future.cancelled(): + self.cancelled[future] = self.futures.pop(future) + index = self.cancelled[future] + elif future.exception(0.1) is not None: + self.failed[future] = self.futures.pop(future) + index = self.failed[future] + else: + self.succeeded[future] = self.futures.pop(future) + index = self.succeeded[future] + if self.callback: + try: + self.callback( + file=self.files[index], + result=future.result(0.1), + exc=future.exception(0.1), + ) + except (CancelledError, TimeoutError): + ... + + # break out downloads as we dont allow failures + if not self.futures: + self.is_running = False + + def start(self): + """submit all files'workers to the executor""" + self.is_running = True + self.futures = { + self.executor.submit( + download_file_worker, + file, + self.mount_point, + self.temp_dir, + self.on_data, + ): index + for index, file in enumerate(self.files) + } + for future in list(self.futures.keys()): + future.add_done_callback(self.notify_completion) + + def shutdown(self, now=True): + if now or self.cancelled or self.failed: + self.executor.shutdown(wait=False, cancel_futures=True) + else: + self.executor.shutdown(wait=True) + self.is_running = False + + +class ProcessingLocalContent(Step): + name = "Processing local contents" + + def run(self, payload, *args, **kwargs) -> int: + mount_point = payload["image"].p3_mounted_on + + # only non-remote Files (plain and local) + for file in payload["config"].non_remote_files: + dest_path = file.mounted_to(mount_point) + try: + ensure_dir(dest_path.parent) + except Exception as exc: + logger.fail_task(str(exc)) + return 1 + + if file.is_plain: + logger.start_task(f"Writing plain text to {file.to}…") + try: + size = dest_path.write_text(file.content) + except Exception as exc: + logger.fail_task(str(exc)) + return 1 + else: + logger.succeed_task(format_size(size)) + continue + + src_path = pathlib.Path(file.url) + + if file.is_direct: + logger.start_task(f"Copying file to {file.to}…") + try: + shutil.copy2(src_path, dest_path) + except Exception as exc: + logger.fail_task(str(exc)) + return 1 + else: + logger.succeed_task(format_size(get_filesize(dest_path))) + continue + + logger.start_task(f"Expanding {self.via} file to {file.to}…") + try: + expand_file(src_path, dest_path, file.via) + except Exception as exc: + logger.fail_task(str(exc)) + return 1 + else: + logger.succeed_task(format_size(get_size_of(dest_path))) + return 0 + + +class DownloadingContent(Step): + name = "Downloading content" + + def run(self, payload, *args, **kwargs) -> int: + mount_point = payload["image"].p3_mounted_on + + nb_remotes = len(payload["config"].remote_files) + + dl_progress = MultiDownloadProgress( + nb_total=nb_remotes, + nb_completed=0, + bytes_total=sum( + [ + file.size if file.size else 0 + for file in payload["config"].remote_files + ] + ), + bytes_received=0, + ) + dl_pb = MultiDownloadProgressBar(dl_progress) + + # dont use multi-downloader for single file or thread==1 + if nb_remotes == 1 or payload["options"].concurrency == 1: + + def on_data(bytes_received: int): + """update both data holder and progress bar on data receival""" + dl_progress.on_data(bytes_received) + dl_pb.update() + + for file in payload["config"].remote_files: + logger.add_task(f"Downloading {file.geturl()} into {file.to}…") + dest_path = file.mounted_to(mount_point) + try: + download_file_worker( + file, + mount_point, + payload["options"].build_dir, + on_data=on_data, + ) + except Exception as exc: + logger.complete_download( + dest_path.name, failed=True, extra=str(exc) + ) + return 1 + else: + dl_progress.nb_completed += 1 + logger.complete_download( + dest_path.name, + format_size(get_size_of(dest_path)), + extra=f"({str(dl_progress)})", + ) + return 0 + + # multi-download with UI refresh on MainThread + def on_completion(file, result: Any, exc: Exception = None): + dl_progress.nb_completed += 1 + dest_path = file.mounted_to(mount_point) + if exc is not None: + logger.debug(f"Failed to download {file.url} into {dest_path}: {exc}") + raise exc + + logger.complete_download( + dest_path.name, + format_size(get_size_of(dest_path)), + extra=f"({str(dl_progress)})", + ) + + downloader = FilesMultiDownloader( + files=payload["config"].remote_files, + mount_point=mount_point, + temp_dir=payload["options"].build_dir.joinpath("dl_remotes"), + concurrency=payload["options"].concurrency, + callback=on_completion, + on_data=dl_progress.on_data, + ) + + logger.add_task( + f"Downloading {nb_remotes} files " + f"totaling {format_size(dl_progress.bytes_total)}…", + f"using {min([nb_remotes, downloader.executor._max_workers])} workers", + ) + + try: + downloader.start() + while downloader.is_running: + dl_pb.update() + except KeyboardInterrupt: + downloader.shutdown(now=True) + # TODO: we should actually interrupt downloads here + raise + else: + downloader.shutdown() + finally: + dl_pb.update() + dl_pb.finish() + return 0 diff --git a/src/image_creator/steps/image.py b/src/image_creator/steps/image.py new file mode 100644 index 0000000..f0cac9b --- /dev/null +++ b/src/image_creator/steps/image.py @@ -0,0 +1,143 @@ +from image_creator.constants import logger +from image_creator.steps import Step +from image_creator.utils.image import Image +from image_creator.utils.misc import format_size + + +class ResizingImage(Step): + name = "Resizing image" + + def run(self, payload, *args, **kwargs) -> int: + logger.start_task("Checking image size…") + payload["image"] = Image(payload["options"].output_path) + try: + size = payload["image"].get_size() + except Exception as exc: + logger.fail_task(str(exc)) + return 1 + else: + logger.succeed_task(format_size(size)) + + logger.start_task(f"Resizing image to {payload['output_size']}b…") + try: + payload["image"].resize(payload["output_size"]) + except Exception as exc: + logger.fail_task(str(exc)) + return 1 + else: + logger.succeed_task(format_size(payload["image"].get_size())) + + logger.start_task("Getting a loop device…") + try: + loop_dev = payload["image"].assign_loopdev() + except Exception as exc: + logger.fail_task(str(exc)) + return 1 + else: + logger.succeed_task(loop_dev) + + logger.start_task(f"Attaching image to {loop_dev}…") + try: + payload["image"].attach() + except Exception as exc: + logger.fail_task(str(exc)) + return 1 + else: + logger.succeed_task() + + logger.start_task(f"Resizing third partition of {loop_dev}…") + try: + payload["image"].resize_last_part() + except Exception as exc: + logger.fail_task(str(exc)) + return 1 + else: + logger.succeed_task() + + return 0 + + def cleanup(self, payload): + if payload.get("image"): + payload["image"].detach() + + +class MountingDataPart(Step): + name = "Mounting data partition" + + def run(self, payload, *args, **kwargs) -> int: + logger.start_task(f"Mouting {payload['image'].loop_dev}p3…") + try: + mounted_on = payload["image"].mount_p3() + except Exception as exc: + logger.fail_task(str(exc)) + return 1 + else: + logger.succeed_task(mounted_on) + return 0 + + def cleanup(self, payload): + if payload.get("image"): + payload["image"].unmount_p3() + + +class UnmountingDataPart(Step): + name = "Unmounting data partition" + + def run(self, payload, *args, **kwargs) -> int: + logger.start_task(f"Unmouting {payload['image'].p3_mounted_on}…") + try: + payload["image"].unmount_p3() + except Exception as exc: + logger.fail_task(str(exc)) + return 1 + else: + logger.succeed_task() + return 0 + + +class UnmountingBootPart(Step): + name = "Unmounting boot partition" + + def run(self, payload, *args, **kwargs) -> int: + logger.start_task(f"Unmouting {payload['image'].p1_mounted_on}…") + try: + payload["image"].unmount_p1() + except Exception as exc: + logger.fail_task(str(exc)) + return 1 + else: + logger.succeed_task() + return 0 + + +class MountingBootPart(Step): + name = "Mounting boot partition" + + def run(self, payload, *args, **kwargs) -> int: + logger.start_task(f"Mouting {payload['image'].loop_dev}p1…") + try: + mounted_on = payload["image"].mount_p1() + except Exception as exc: + logger.fail_task(str(exc)) + return 1 + else: + logger.succeed_task(mounted_on) + return 0 + + def cleanup(self, payload): + if payload.get("image"): + payload["image"].unmount_p1() + + +class DetachingImage(Step): + name = "Detaching Image" + + def run(self, payload, *args, **kwargs) -> int: + logger.start_task(f"Detach image from {payload['image'].loop_dev}") + if not payload["image"].detach(): + logger.fail_task(f"{payload['image']} not detached!") + return 1 + else: + logger.succeed_task() + + return 0 diff --git a/src/image_creator/steps/oci_images.py b/src/image_creator/steps/oci_images.py new file mode 100644 index 0000000..0f2dd3e --- /dev/null +++ b/src/image_creator/steps/oci_images.py @@ -0,0 +1,37 @@ +from image_creator.constants import logger +from image_creator.steps import Step +from image_creator.utils.misc import format_size, get_filesize, rmtree +from image_creator.utils.oci_images import download_image + + +class DownloadingOCIImages(Step): + name = "Downloading OCI Images" + + def run(self, payload, *args, **kwargs) -> int: + logger.start_task("Creating OCI Images placeholder…") + mount_point = payload["image"].p3_mounted_on + images_dir = mount_point.joinpath("images") + build_dir = payload["options"].build_dir.joinpath("oci_export") + + try: + images_dir.mkdir(exist_ok=True, parents=True) + except Exception as exc: + logger.fail_task(str(exc)) + return 1 + else: + logger.succeed_task(images_dir) + + for image in payload["config"].oci_images: + target = images_dir.joinpath(f"{image.fs_name}.tar") + logger.add_task( + f"Downloading OCI Image to {target.relative_to(mount_point)}…" + ) + try: + download_image(image=image, dest=target, build_dir=build_dir) + except Exception as exc: + logger.fail_task(str(exc)) + rmtree(build_dir) + return 1 + else: + logger.complete_download(target.name, format_size(get_filesize(target))) + return 0 diff --git a/src/image_creator/steps/sizes.py b/src/image_creator/steps/sizes.py new file mode 100644 index 0000000..2422f80 --- /dev/null +++ b/src/image_creator/steps/sizes.py @@ -0,0 +1,24 @@ +from image_creator.constants import logger +from image_creator.steps import Step +from image_creator.utils.misc import format_size, parse_size + + +class ComputeSizes(Step): + name = "Compute sizes…" + + def run(self, payload, *args, **kwargs) -> int: + + payload["output_size"] = parse_size(payload["config"].dig("output.size")) + logger.add_task("Image size:", f"{format_size(payload['output_size'])}") + + # TODO: we should do more + # - compute size of all content + # - display cumulative size of content + # - display suggested minimum SD size + # - use that size as image size if set to `auto` + # - fail if `auto` but some files dont report size + # - check disk space on output_path.parent + # - failed it it wont allow the base + resize + # - calc how much space will be needed in the build-dir (non-direct DL) + # - fail if build-dir wont allow it + return 0 diff --git a/src/image_creator/utils/__init__.py b/src/image_creator/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/image_creator/utils/download.py b/src/image_creator/utils/download.py new file mode 100644 index 0000000..66b21c5 --- /dev/null +++ b/src/image_creator/utils/download.py @@ -0,0 +1,109 @@ +import io +import pathlib +from typing import Callable, Dict, Optional, Union + +import requests + +session = requests.Session() +# basic urllib retry mechanism. +# Sleep (seconds): {backoff factor} * (2 ** ({number of total retries} - 1)) +# https://docs.descarteslabs.com/_modules/urllib3/util/retry.html +retries = requests.packages.urllib3.util.retry.Retry( + total=10, # Total number of retries to allow. Takes precedence over other counts. + connect=5, # How many connection-related errors to retry on + read=5, # How many times to retry on read errors + redirect=20, # How many redirects to perform. (to avoid infinite redirect loops) + status=3, # How many times to retry on bad status codes + other=0, # How many times to retry on other errors + allowed_methods=False, # Set of HTTP verbs that we should retry on (False is all) + status_forcelist=[ + 413, + 429, + 500, + 502, + 503, + 504, + ], # Set of integer HTTP status we should force a retry on + backoff_factor=30, # backoff factor to apply between attempts after the second try, + raise_on_redirect=False, # raise MaxRetryError instead of 3xx response + raise_on_status=False, # raise on Bad Status or response + respect_retry_after_header=True, # respect Retry-After header (status_forcelist) +) +retries.DEFAULT_BACKOFF_MAX = 30 * 60 # allow up-to 30mn backoff (default 2mn) +session.mount("http", requests.adapters.HTTPAdapter(max_retries=retries)) + + +def get_online_rsc_size(url: str) -> int: + """size (Content-Length) from url if specified, -1 otherwise (-2 on errors)""" + try: + resp = requests.head(url, allow_redirects=True, timeout=60) + # some servers dont offer HEAD + if resp.status_code != 200: + resp = requests.get( + url, + allow_redirects=True, + timeout=60, + stream=True, + headers={"Accept-Encoding": "identity"}, + ) + resp.raise_for_status() + return int(resp.headers.get("Content-Length") or -1) + except Exception: + return -2 + + +def read_text_from(url: str) -> str: + """Text content from an URL""" + resp = session.get(url) + resp.raise_for_status() + return resp.text + + +def download_file( + url: str, + to: Union[pathlib.Path, io.BytesIO], + block_size: Optional[int] = 1024, + proxies: Optional[dict] = None, + only_first_block: Optional[bool] = False, + headers: Optional[Dict[str, str]] = None, + on_data: Optional[Callable] = None, +) -> Union[int, requests.structures.CaseInsensitiveDict]: + """Stream data from a URL to either a BytesIO object or a file + Arguments - + fpath - Path or BytesIO to write data into + block_size - Size of each chunk of data read in one iteration + proxies - A dict of proxies to be used + https://requests.readthedocs.io/en/master/user/advanced/#proxies + only_first_block - Whether to download only one (first) block + + Returns the total number of bytes downloaded and the response headers""" + + resp = session.get( + url, + stream=True, + proxies=proxies, + headers=headers, + ) + resp.raise_for_status() + + total_downloaded = 0 + if isinstance(to, pathlib.Path): + fp = open(to, "wb") + + for data in resp.iter_content(block_size): + nb_received = len(data) + total_downloaded += nb_received + fp.write(data) + + if on_data: + on_data(nb_received) + + # stop downloading/reading if we're just testing first block + if only_first_block: + break + + if isinstance(to, pathlib.Path): + fp.close() + else: + fp.seek(0) + return total_downloaded, resp.headers diff --git a/src/image_creator/utils/image.py b/src/image_creator/utils/image.py new file mode 100644 index 0000000..e230b8a --- /dev/null +++ b/src/image_creator/utils/image.py @@ -0,0 +1,298 @@ +import json +import pathlib +import re +import subprocess +import tempfile +from typing import Optional + +from image_creator.utils.misc import get_environ, rmtree + + +def get_image_size(fpath: pathlib.Path) -> int: + """Size in bytes of the virtual device in image""" + virtsize_re = re.compile( + r"^virtual size: ([0-9\.\sa-zA-Z]+) \((?P\d+) bytes\)" + ) + for line in subprocess.run( + ["/usr/bin/env", "qemu-img", "info", "-f", "raw", str(fpath)], + check=True, + capture_output=True, + text=True, + env=get_environ(), + ).stdout.splitlines(): + match = virtsize_re.match(line) + if match: + return int(match.groupdict().get("size")) + + +def resize_image(fpath: pathlib.Path, size: int): + """Resize virtual device in image (bytes)""" + subprocess.run( + ["/usr/bin/env", "qemu-img", "resize", "-f", "raw", fpath, str(size)], + check=True, + capture_output=True, + text=True, + env=get_environ(), + ) + + +def get_loopdev() -> str: + """free loop-device path ready to ease""" + return subprocess.run( + ["/usr/bin/env", "losetup", "-f"], + check=True, + capture_output=True, + text=True, + env=get_environ(), + ).stdout.strip() + + +def is_loopdev_free(loop_dev: str): + """whether a loop-device (/dev/loopX) is not already attached""" + devices = json.loads( + subprocess.run( + ["/usr/bin/env", "losetup", "--json"], + check=True, + capture_output=True, + text=True, + env=get_environ(), + ).stdout.strip() + )["loopdevices"] + return loop_dev not in [device["name"] for device in devices] + + +def attach_to_device(img_fpath: pathlib.Path, loop_dev: str): + """attach a device image to a loop-device""" + subprocess.run( + ["/usr/bin/env", "losetup", "--partscan", loop_dev, str(img_fpath)], + check=True, + capture_output=True, + text=True, + env=get_environ(), + ) + + +def detach_device(loop_dev: str, failsafe: Optional[bool] = False) -> bool: + """whether detaching this loop-device succeeded""" + return ( + subprocess.run( + ["/usr/bin/env", "losetup", "--detach", loop_dev], + check=not failsafe, + capture_output=True, + text=True, + env=get_environ(), + ).returncode + == 0 + ) + + +def get_device_sectors(dev_path: str) -> int: + """number of sectors composing this device""" + summary_re = re.compile( + rf"^Disk {dev_path}: (?P[\d\.\s]+ [KMGP]iB+), " + r"(?P\d+) bytes, (?P\d+) sectors$" + ) + line = subprocess.run( + ["/usr/bin/env", "fdisk", "--list", dev_path], + check=True, + capture_output=True, + text=True, + env=get_environ(), + ).stdout.splitlines()[0] + match = summary_re.match(line) + if match: + return int(match.groupdict().get("sectors")) + raise ValueError(f"Unable to get nb of sectors for {dev_path}") + + +def get_thirdpart_start_sector(dev_path) -> int: + """Start sector number of third partition of device""" + part_re = re.compile( + rf"{dev_path}p3 (?P\d+) (?P\d+) (?P\d+) .+$" + ) + line = subprocess.run( + ["/usr/bin/env", "fdisk", "--list", dev_path], + check=True, + capture_output=True, + text=True, + env=get_environ(), + ).stdout.splitlines()[-1] + match = part_re.match(line) + if match: + return int(match.groupdict().get("start")) + raise ValueError(f"Unable to get start sector for {dev_path}p3") + + +def resize_third_partition(dev_path: str): + """recreate third partition of a device and its (ext4!) filesystem""" + nb_sectors = get_device_sectors(dev_path) + start_sector = get_thirdpart_start_sector(dev_path) + end_sector = nb_sectors - 1 + + # delete 3rd part and recreate from same sector until end of device + commands = ["d", "3", "n", "p", "3", str(start_sector), str(end_sector), "N", "w"] + subprocess.run( + ["/usr/bin/env", "fdisk", dev_path], + check=True, + input="\n".join(commands), + capture_output=True, + text=True, + env=get_environ(), + ) + + # check fs on 3rd part + subprocess.run( + ["/usr/bin/env", "e2fsck", "-p", f"{dev_path}p3"], + check=True, + capture_output=True, + text=True, + env=get_environ(), + ) + + # resize fs on 3rd part + subprocess.run( + ["/usr/bin/env", "resize2fs", f"{dev_path}p3"], + check=True, + capture_output=True, + text=True, + env=get_environ(), + ) + + +def mount_on( + dev_path: str, mount_point: pathlib.Path, filesystem: Optional[str] +) -> bool: + """whether mounting device onto mount point succeeded""" + commands = ["/usr/bin/env", "mount"] + if filesystem: + commands += ["-t", filesystem] + commands += [dev_path, str(mount_point)] + return ( + subprocess.run( + commands, + capture_output=True, + text=True, + env=get_environ(), + ).returncode + == 0 + ) + + +def unmount(mount_point: pathlib.Path) -> bool: + """whether unmounting mount-point succeeded""" + return ( + subprocess.run( + ["/usr/bin/env", "umount", str(mount_point)], + capture_output=True, + text=True, + env=get_environ(), + ).returncode + == 0 + ) + + +class Image: + """File-backed Image that can be attached/resized/mounted""" + + def __init__(self, fpath: pathlib.Path, mount_in: Optional[pathlib.Path] = None): + # ensure image is readable + with open(fpath, "rb") as fh: + fh.read(1024) + self.fpath = fpath + self.loop_dev = None + self.p1_mounted_on = None + self.p3_mounted_on = None + self.mount_in = mount_in + + @property + def is_mounted(self) -> bool: + return bool(self.p3_mounted_on) or bool(self.p1_mounted_on) + + @property + def is_attached(self) -> bool: + return bool(self.loop_dev) + + def get_size(self) -> int: + """virtual device size""" + return get_image_size(self.fpath) + + def assign_loopdev(self) -> str: + """find a free loop device we'll use""" + self.loop_dev = get_loopdev() + return self.loop_dev + + def attach(self): + """attach image to loop device""" + if not self.loop_dev or not is_loopdev_free(self.loop_dev): + detach_device(self.loop_dev, failsafe=True) + self.loop_dev = None + self.assign_loopdev() + attach_to_device(self.fpath, self.loop_dev) + + def detach(self): + """detach loop-device""" + if self.is_mounted: + self.unmount_all() + if self.loop_dev and detach_device(self.loop_dev): + self.loop_dev = None + return True + return False + + def resize(self, to: int): + """resize virtual device inside image (expand only)""" + resize_image(self.fpath, size=to) + + def resize_last_part(self): + """resize 3rd partition and filesystem to use all remaining space""" + resize_third_partition(self.loop_dev) + + def mount_p1(self) -> pathlib.Path: + """mount first (boot) partition""" + self.mount_part(1) + + def mount_p3(self) -> pathlib.Path: + """mount third (data) partition""" + self.mount_part(3) + + def unmount_p1(self) -> pathlib.Path: + """unmount first (boot) partition""" + self.unmount_part(1) + + def unmount_p3(self) -> pathlib.Path: + """unmount third (data) partition""" + self.unmount_part(3) + + def mount_part(self, part_num: int) -> pathlib.Path: + """path to mounted specific partition""" + mount_point = pathlib.Path( + tempfile.mkdtemp(dir=self.mount_in, prefix=f"part{part_num}_") + ) + fs = "msdos" if part_num == 1 else "ext4" + if mount_on(f"{self.loop_dev}p{part_num}", mount_point, fs): + setattr(self, f"p{part_num}_mounted_on", mount_point) + else: + raise IOError( + f"Unable to mount {self.loop_dev}p{part_num} on {mount_point}" + ) + + def unmount_part(self, part_num: int): + """unmount specific partition""" + mount_point = getattr(self, f"p{part_num}_mounted_on") + if not mount_point: + return + if unmount(mount_point): + setattr(self, f"p{part_num}_mounted_on", None) + rmtree(mount_point) + else: + raise IOError(f"Unable to unmount p{part_num} at {mount_point}") + + def unmount_all(self): + """failsafely unmount all partitions we would have mounted""" + try: + self.unmount_p1() + except Exception: + ... + try: + self.unmount_p3() + except Exception: + ... diff --git a/src/image_creator/utils/misc.py b/src/image_creator/utils/misc.py new file mode 100644 index 0000000..915a25c --- /dev/null +++ b/src/image_creator/utils/misc.py @@ -0,0 +1,94 @@ +import lzma +import os +import pathlib +import shutil +import tarfile +import zipfile +from typing import Dict + +import humanfriendly + + +def format_size(size: int) -> str: + """human-readable representation of a size in bytes""" + return humanfriendly.format_size(size, binary=True) + + +def parse_size(size: str) -> int: + """size in bytes of a human-readable size representation""" + return humanfriendly.parse_size(size) + + +def get_filesize(fpath: pathlib.Path) -> int: + """size in bytes of a local file path""" + return fpath.stat().st_size + + +def get_dirsize(fpath: pathlib.Path) -> int: + """size in bytes of a local directory""" + if not fpath.exists(): + raise FileNotFoundError(fpath) + if fpath.is_file(): + raise IOError(f"{fpath} is a file") + return sum(f.stat().st_size for f in fpath.rglob("**/*") if f.is_file()) + + +def get_size_of(fpath: pathlib.Path) -> int: + """size in bytes of a local file or directory""" + if not fpath.exists(): + raise FileNotFoundError(fpath) + if fpath.is_file(): + return get_filesize(fpath) + return get_dirsize(fpath) + + +def rmtree(fpath: pathlib.Path): + """recursively remove an entire folder (rm -rf)""" + shutil.rmtree(fpath, ignore_errors=True) + + +def ensure_dir(fpath: pathlib.Path): + """recursively creating a folder (mkdir -p)""" + fpath.mkdir(parents=True, exist_ok=True) + + +def get_environ() -> Dict[str, str]: + """current environment variable with langs set to C to control cli output""" + environ = os.environ.copy() + environ.update({"LANG": "C", "LC_ALL": "C"}) + return environ + + +def extract_xz_image(src: pathlib.Path, dest: pathlib.Path): + """Extract compressed (lzma via xz compress) image file""" + buff_size = parse_size("16MiB") + buffer = b"" + with lzma.open(src, "rb") as reader, open(dest, "wb") as writer: + buffer = reader.read(buff_size) + while buffer: + writer.write(buffer) + buffer = reader.read(buff_size) + + +def expand_file(src: pathlib.Path, method: str, dest: pathlib.Path): + """Expand into dest failing should any member to-be written outside dest""" + if method not in shutil._UNPACK_FORMATS.keys(): + raise NotImplementedError(f"Cannot expand `{method}`") + + # raise on unauthorized filenames instead of ignoring (zip) or accepting (tar) + names = [] + if method == "zip": + with zipfile.ZipFile(src, "r") as zh: + names = zh.namelist() + elif method == "tar" or method.endswith("tar"): + with tarfile.Tarfile(src, "r") as th: + names = th.getnames() + for name in names: + path = pathlib.Path(name) + if path.root == "/": + raise IOError(f"{method} file contains member with absolute path: {name}") + path = dest.joinpath(name).resolve() + if not path.is_relative_to(dest): + raise IOError(f"{method} file contains out-of-bound member path: {name}") + + return shutil.unpack_archive(src, dest, method) diff --git a/src/image_creator/utils/oci_images.py b/src/image_creator/utils/oci_images.py new file mode 100644 index 0000000..c12a62d --- /dev/null +++ b/src/image_creator/utils/oci_images.py @@ -0,0 +1,25 @@ +import logging +import pathlib + +from docker_export import Image, Platform, RegistryAuth, export, get_layers_manifest +from image_creator.constants import logger + +logging.getLogger("docker_export").setLevel(logging.WARNING) +platform = Platform.parse("linux/arm64/v8") + + +def image_exists(image: Image) -> bool: + """whether image exists on the registry""" + auth = RegistryAuth.init(image) + auth.authenticate() + try: + get_layers_manifest(image=image, platform=platform, auth=auth) + except Exception as exc: + logger.exception(exc) + return False + return True + + +def download_image(image: Image, dest: pathlib.Path, build_dir: pathlib.Path): + """download image into a tar file at dest""" + export(image=image, platform=platform, to=dest, build_dir=build_dir) diff --git a/src/image_creator/utils/requirements.py b/src/image_creator/utils/requirements.py new file mode 100644 index 0000000..dd1c4fe --- /dev/null +++ b/src/image_creator/utils/requirements.py @@ -0,0 +1,69 @@ +import os +import re +import subprocess +from typing import List, Tuple + +help_text = """ +Requirements +------------ + +kernel features: + - `loop` must be enabled in your kernel or as a module + if running inside a docker-container: + - same loop feature applies to host's kernel + - container must be run with --privileged + - `ext4` filesystem (most likely enabled in-kernel) + check whether you need t + +tools: + - losetup (mount) + - fdisk (fdisk) + - resize2fs (e2fsprogs) + - mount (mount) + - umount (mount) + - qemu-img (qemu-utils) + +Sample setup (debian) +sudo modprobe --first-time loop +sudo modprobe --first-time ext4 +sudo apt-get install --no-install-recommends mount fdisk e2fsprogs qemu-utils +""" + + +def is_root() -> bool: + """whether running as root""" + return os.getuid() == 0 + + +def has_ext4_support() -> bool: + """whether ext4 filesystem is enabled""" + with open("/proc/filesystems", "r") as fh: + for line in fh.readlines(): + if re.match(r"\s*ext4\s?$", line.strip()): + return True + return False + + +def has_all_binaries() -> Tuple[bool, List]: + """whether all required binaries are present, with list of missing ones""" + missing_bins = [] + for binary in ("losetup", "fdisk", "resize2fs", "mount", "umount", "qemu-img"): + try: + if ( + subprocess.run(["/usr/bin/env", binary], capture_output=True).returncode + == 127 + ): + missing_bins.append(binary) + except Exception: + missing_bins.append(binary) + return not missing_bins, missing_bins + + +def has_loop_device() -> bool: + """whether requesting a loop-device is possible""" + return ( + subprocess.run( + ["/usr/bin/env", "losetup", "-f"], capture_output=True + ).returncode + == 0 + )