archiver.py

#!/usr/bin/env python3
import math
import os
import tarfile
import zipfile
from datetime import datetime
from hashlib import file_digest
from io import BytesIO
from pathlib import Path
from subprocess import run


class Disk:
    device: Path

    def __init__(self, device: Path):
        self.device = device

    # noinspection SpellCheckingInspection
    def get_device_size_bytes(self) -> int:
        # fmt: off
        command = ["lsblk", str(self.device), "--output", "SIZE", "--bytes", "--noheadings", "--nodeps"]
        # fmt: on
        output = run(command, capture_output=True, check=True)
        return int(output.stdout.decode())

    def write_disk(self, image: bytes):
        disk_size = self.get_device_size_bytes()
        if len(image) > disk_size:
            print("Error: image size exceeds disk size")
            exit(1)
        print(self.device)
        copies = max(1, math.floor(disk_size / len(image)))
        image *= copies
        print("disk size: %d bytes" % disk_size)
        print("disk usage: %0.2f%%" % ((len(image) / disk_size) * 100))
        print("image copies: %d" % copies)
        print("writing image to disk...")
        start = datetime.now()
        self.device.write_bytes(image)
        end = datetime.now()
        delta = end - start
        print("write time:", round(delta.total_seconds(), 2), "seconds")
        print("write speed:", round(len(image) / delta.total_seconds() / 1024, 2), "KiB/s")


class ArchiveBuilder:
    base_dir: Path
    files: list[Path]
    squashfs_options: list[str]

    def __init__(self, base_dir: Path, squashfs_options: list[str]):
        self.base_dir = base_dir
        self.squashfs_options = squashfs_options
        self.files = sorted(
            Path(root) / file
            for root, _, files in self.base_dir.walk()
            for file in files
        )

    def build_checksum_file(self, digest: str) -> str:
        lines: list[str] = []
        for file in self.files:
            with file.open("rb") as file_obj:
                file_hash = file_digest(file_obj, digest).hexdigest()
                file_name = file.relative_to(self.base_dir)
                lines.append(f"{file_hash}  {file_name}")
        return "\n".join(lines) + "\n"

    def to_zip_archive(self) -> bytes:
        zip_archive = BytesIO()

        with zipfile.ZipFile(file=zip_archive, mode="w") as zip_file:
            zip_file: zipfile.ZipFile
            zip_file.compression = zipfile.ZIP_DEFLATED
            zip_file.compresslevel = 9
            zip_file.allowZip64 = False
            zip_file.strict_timestamps = True
            for file in self.files:
                zip_file.write(file, file.relative_to(self.base_dir))
            zip_file.writestr("SHA1SUMS", self.build_checksum_file("sha1"))
            zip_file.writestr("SHA256SUMS", self.build_checksum_file("sha256"))

        return zip_archive.getvalue()

    def to_tar_archive(self) -> bytes:
        tar_archive = BytesIO()

        def emit_checksum_file(archive: tarfile.TarFile, digest: str):
            content = self.build_checksum_file(digest)
            tar_info = tarfile.TarInfo(digest.upper() + "SUMS")
            tar_info.size = len(content)
            archive.addfile(tar_info, BytesIO(content.encode("utf-8")))

        with tarfile.open(fileobj=tar_archive, mode="w") as tar_file:
            tar_file: tarfile.TarFile
            for file in self.files:
                tar_file.add(file, arcname=str(file.relative_to(self.base_dir)))
            emit_checksum_file(tar_file, digest="sha1")
            emit_checksum_file(tar_file, digest="sha256")

        return tar_archive.getvalue()

    def write_zip_archive(self, zip_file: Path):
        zip_file.unlink(missing_ok=True)
        zip_file.write_bytes(self.to_zip_archive())
        print("zip archive:", zip_file.stat().st_size, "bytes")

    def write_squashfs_image(self, image: Path):
        image_logs = image.with_suffix(f"{image.suffix}.log")
        image_info = image.with_suffix(f"{image.suffix}.info")
        image.unlink(missing_ok=True)
        image_logs.unlink(missing_ok=True)
        image_info.unlink(missing_ok=True)
        with image_logs.open("w") as stdout:
            cwd = image.parent.resolve()
            # noinspection SpellCheckingInspection
            cmd = ["mksquashfs", "-", image.name, "-info-file", image_info.name, "-tar", *self.squashfs_options]
            run(cmd, input=self.to_tar_archive(), cwd=cwd, stdout=stdout, stderr=stdout, check=True)
        print("squashfs image:", image.stat().st_size, "bytes")


def main(base_dir: Path, archive_dir: str, name: str, disk_pattern: str):
    zip_archive = base_dir / f"{name}.zip"
    squashfs_file = base_dir / f"{name}.squashfs"
    archive = ArchiveBuilder(base_dir.joinpath(archive_dir), squashfs_options=[
        # fmt: off
        "-noappend",
        "-no-progress", "-no-xattrs", "-no-exports", "-no-sparse", "-no-tailends",
        "-all-root", "-force-file-mode", "400", "-force-dir-mode", "500",
        "-comp", "gzip", "-Xcompression-level", "9",
        "-exit-on-error",
        # fmt: on
    ])
    archive.write_zip_archive(zip_archive)
    archive.write_squashfs_image(squashfs_file)
    for disk in Path("/dev/disk/by-id/").glob(disk_pattern):
        Disk(disk).write_disk(squashfs_file.read_bytes())


if __name__ == "__main__":
    main(
        base_dir=Path(__file__).parent.resolve(),
        archive_dir="files",
        name=os.environ.get("ARCHIVE_NAME", "mini-disk-archive"),
        disk_pattern=os.environ.get("ARCHIVE_DISK_PATTERN", "usb-wch.cn_USB-Mini-Disk*-0:0"),
    )