path: root/modifyiso.py

#!/usr/bin/env python3
import argparse
import logging
import os
import re
import shlex
import shutil
import subprocess
import sys
from time import monotonic, sleep
from typing import Union, Optional


class ModifyIso:
    logger = logging

    def __init__(self, source_iso, temp_dir, chroot_script, target_iso, interactive=False, hostname="vyos",
                 skip_extract=False, skip_chroot=False, skip_iso=False, stdout_only=False):
        self.source_iso: str = os.path.realpath(source_iso)
        self.temp_dir = temp_dir
        self.chroot_script = chroot_script
        self.target_iso = target_iso
        self.interactive = interactive
        self.hostname = hostname
        self.skip_extract = skip_extract
        self.skip_chroot = skip_chroot
        self.skip_iso = skip_iso
        self.stdout_only = stdout_only

        self.mount_dir = os.path.join(self.temp_dir, "mount")
        self.iso_dir = os.path.join(self.temp_dir, "iso")
        self.squashfs_dir = os.path.join(self.temp_dir, "squashfs")
        self.original_working_dir = os.getcwd()

    def run(self):
        if not self.skip_extract:
            self.mount_and_extract()

        if not self.skip_chroot:
            self.execute_chroot()

        if not self.skip_iso:
            self.generate_iso()

    def mount_and_extract(self):
        self.mount_iso(self.source_iso)

        if os.path.exists(self.iso_dir):
            shutil.rmtree(self.iso_dir)
        os.makedirs(self.iso_dir)

        self.logger.info("Copying ISO contents to '%s'" % self.iso_dir)
        self.execute(["/usr/bin/cp", "-a", os.path.join(self.mount_dir, "."), self.iso_dir])

        self.logger.info("Unmounting ISO")
        self.umount(self.mount_dir)

        if os.path.exists(self.squashfs_dir):
            self.umount_all(self.squashfs_dir)
            shutil.rmtree(self.squashfs_dir)

        self.logger.info("Extracting squashfs filesystem")
        os.chdir(self.temp_dir)
        self.execute([
            "/usr/bin/unsquashfs",
            "-d", self.squashfs_dir,
            os.path.join(self.iso_dir, "live/filesystem.squashfs")
        ])

    def execute_chroot(self):
        os.chdir(self.original_working_dir)
        self.logger.info("Preparing chroot")

        apt_sources_path = os.path.join(self.squashfs_dir, "etc/apt/sources.list")
        apt_sources_backup_path = "%s.bak" % apt_sources_path
        if not os.path.exists(apt_sources_backup_path):
            os.rename(apt_sources_path, apt_sources_backup_path)

        with open(apt_sources_backup_path, "r") as source_file:
            lines = []
            for line in source_file:
                url = "http://ftp.debian.org/debian/"
                if "security" in line:
                    url = "http://security.debian.org/debian-security"
                line = re.sub(r"(deb|deb-src)\s+[^\s]+\s+", "\\1 %s " % url, line)
                lines.append(line)

            with open(apt_sources_path, "w") as target_file:
                for line in lines:
                    target_file.write(line)
                    target_file.write("\n")

        resolv_path = os.path.join(self.squashfs_dir, "etc/resolv.conf")
        resolv_backup_path = "%s.bak" % resolv_path
        if not os.path.exists(resolv_backup_path):
            os.rename(resolv_path, resolv_backup_path)
        shutil.copy2("/etc/resolv.conf", resolv_path)

        for directory in ["dev", "proc", "sys"]:
            target = os.path.join(self.squashfs_dir, directory)
            self.execute(["/usr/bin/mount", "--bind", os.path.join("/", directory), target])

        temp_script_path = os.path.join(self.squashfs_dir, "custom.sh")
        chroot_temp_script_path = temp_script_path[len(self.squashfs_dir):]
        if self.chroot_script:
            shutil.copy2(self.chroot_script, temp_script_path)
            os.chmod(temp_script_path, 0o755)

        try:
            self.logger.info("Executing chroot '%s'" % self.squashfs_dir)

            if self.hostname:
                import unshare
                unshare.unshare(unshare.CLONE_NEWUTS)
                self.execute(["/usr/bin/hostname", self.hostname])

            if self.chroot_script:
                self.execute(["/usr/sbin/chroot", self.squashfs_dir, "/bin/sh", "-c", chroot_temp_script_path])

            if self.interactive:
                self.execute(["/usr/sbin/chroot", self.squashfs_dir, "/bin/sh"], stdin=sys.stdin)

            self.logger.info("Cleaning chroot")

            if os.path.exists(temp_script_path):
                os.remove(temp_script_path)

            os.remove(apt_sources_path)
            os.rename(apt_sources_backup_path, apt_sources_path)

            os.remove(resolv_path)
            os.rename(resolv_backup_path, resolv_path)

        finally:
            for directory in ["dev", "proc", "sys"]:
                target = os.path.join(self.squashfs_dir, directory)
                self.execute(["/usr/bin/umount", "-l", target], valid_codes=[0, 32])

    def generate_iso(self):
        self.logger.info("Generating squashfs")

        squashfs_path = os.path.join(self.iso_dir, "live/filesystem.squashfs")
        if os.path.exists(squashfs_path):
            os.remove(squashfs_path)

        self.execute(["/usr/bin/mksquashfs", self.squashfs_dir, squashfs_path, "-comp", "xz"])

        self.logger.info("Generating checksums")
        self.generate_checksums(self.iso_dir)

        self.logger.info("Generating ISO")
        target_iso_path = self.target_iso
        if target_iso_path == "auto":
            parts = os.path.splitext(self.source_iso)
            target_iso_path = "%s-custom%s" % parts

        if os.path.exists(target_iso_path):
            os.remove(target_iso_path)

        os.chdir(self.original_working_dir)
        command = [
            "/usr/bin/xorriso",
            "-as", "mkisofs",
        ]
        command.extend(self.get_xorriso_arguments(self.source_iso))
        command.extend([
            "-output", target_iso_path,
            self.iso_dir,
        ])
        self.execute(command)

        self.verify_iso(target_iso_path)

        self.logger.info("ISO '%s' generated successfully" % target_iso_path)

    def get_xorriso_arguments(self, iso_path):
        command = [
            "/usr/bin/xorriso",
            "-indev", iso_path,
            "-report_el_torito", "as_mkisofs",
        ]
        output = self.execute(command, stdout=subprocess.PIPE)
        return shlex.split(str(output))

    def generate_checksums(self, directory):
        os.chdir(directory)
        excluded_paths = [
            "./sha256sum.txt",
            "./isolinux",
        ]
        with open("sha256sum.txt", "w") as file:
            for parent, directories, files in os.walk("."):
                for file_name in files:
                    path = os.path.join(parent, file_name)

                    excluded = False
                    for pattern in excluded_paths:
                        if path.startswith(pattern):
                            excluded = True
                            break

                    if excluded:
                        continue

                    output = self.execute(["/usr/bin/sha256sum", path], stdout=subprocess.PIPE)
                    file.write(str(output).strip())
                    file.write("\n")

    def verify_iso(self, iso_path):
        self.logger.info("Verifying ISO")

        self.mount_iso(iso_path, silent=True)

        os.chdir(self.iso_dir)
        self.execute(["/usr/bin/sha256sum", "-c", "sha256sum.txt"], stdout=subprocess.DEVNULL)

        self.umount(self.iso_dir)

        self.logger.info("ISO verified successfully")

    def mount_iso(self, iso_path, silent=False):
        if not os.path.exists(iso_path):
            raise ErrorException("Source ISO '%s' doesn't exist" % iso_path)

        if os.path.exists(self.mount_dir):
            self.umount(self.mount_dir)
        else:
            os.makedirs(self.mount_dir)

        if not silent:
            self.logger.info("Mounting ISO '%s' to '%s'" % (iso_path, self.mount_dir))

        self.execute(["/usr/bin/mount", "-o", "loop,ro", iso_path, self.mount_dir])

    def umount(self, mount):
        self.execute(["/usr/bin/umount", mount], valid_codes=[0, 32], stdout=subprocess.DEVNULL)

    def umount_all(self, target_dir):
        target_dir = os.path.realpath(target_dir)

        related_mounts = []
        with open("/proc/mounts", "r") as file:
            for line in file:
                parts = line.split()
                if len(parts) >= 2:
                    mount_point = os.path.realpath(parts[1].encode("utf-8").decode("unicode_escape"))
                    if mount_point == target_dir or mount_point.startswith(target_dir + os.sep):
                        related_mounts.append(mount_point)

        related_mounts.sort(key=lambda path: path.count(os.sep), reverse=True)

        for mount_point in related_mounts:
            self.umount(mount_point)

    def execute(self, command, valid_codes: Optional[Union[int, list]] = 0, timeout=0, **kwargs) -> Optional[str]:
        if "stdout" not in kwargs:
            kwargs.update({
                "stdout": sys.stdout,
                "stderr": sys.stdout if self.stdout_only else sys.stderr,
                "stdin": sys.stdin,
            })

        if "stderr" not in kwargs:
            kwargs["stderr"] = subprocess.DEVNULL

        process = subprocess.Popen(command, **kwargs)

        deadline = monotonic() + timeout
        while process.poll() is None:
            if timeout > 0 and monotonic() > deadline:
                process.kill()
                raise TimeoutError("Command '%s' timed out" % command)
            sleep(0.100)

        if valid_codes is not None:
            if not isinstance(valid_codes, list):
                valid_codes = [valid_codes]

        if process.returncode not in valid_codes:
            raise ProcessException(command, process.returncode, process.stdout, process.stderr)

        if process.stdout is None:
            return None
        return process.stdout.read().decode("utf-8")


class ProcessException(Exception):
    def __init__(self, command, returncode, stdout=None, stderr=None):
        self.command = command
        self.returncode = returncode
        self.stdout = stdout
        self.stderr = stderr

    def __str__(self):
        message = "Command '%s' failed with code %s" % (self.command, self.returncode)
        if self.stdout:
            message += ", stdout: %s" % self.stdout.read().decode("utf-8")
        if self.stderr:
            message += ", stderr: %s" % self.stderr.read().decode("utf-8")
        return message


def main():
    project_dir = os.path.realpath(os.path.dirname(__file__))

    parser = argparse.ArgumentParser()
    parser.add_argument("source_iso")
    parser.add_argument("--temp-dir", default=os.path.join(project_dir, "build"))
    parser.add_argument("--chroot-script")
    parser.add_argument("--target-iso", default="auto")
    parser.add_argument("--interactive", action="store_true")
    parser.add_argument("--hostname")
    parser.add_argument("--skip-extract", action="store_true")
    parser.add_argument("--skip-chroot", action="store_true")
    parser.add_argument("--skip-iso", action="store_true")

    args = parser.parse_args()
    values = vars(args)

    ModifyIso(**values).run()


class ErrorException(Exception):
    pass


class LessThanLevelFilter(logging.Filter):
    def __init__(self, exclusive_maximum, name="LessThanLevelFilter"):
        super(LessThanLevelFilter, self).__init__(name)
        self.maximum_level = exclusive_maximum

    def filter(self, record):
        return 1 if record.levelno < self.maximum_level else 0


def setup_logging():
    logger = logging.getLogger()
    logger.setLevel(logging.INFO)
    formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")

    stderr_level = logging.WARNING

    stdout_handler = logging.StreamHandler(sys.stdout)
    stdout_handler.setLevel(logging.INFO)
    stdout_handler.addFilter(LessThanLevelFilter(stderr_level))
    stdout_handler.setFormatter(formatter)
    logger.addHandler(stdout_handler)

    stderr_handler = logging.StreamHandler(sys.stderr)
    stderr_handler.setLevel(stderr_level)
    stderr_handler.setFormatter(formatter)
    logger.addHandler(stderr_handler)


if __name__ == "__main__":
    setup_logging()

    try:
        main()

    except ErrorException as e:
        logging.error(e)
        exit(1)
    except KeyboardInterrupt:
        exit(1)
    except Exception as e:
        logging.exception(e)
        exit(1)