#!/usr/bin/python3
import argparse
import pathlib
import subprocess
import tempfile
import pathlib

__author__ = "Trent W. Buck"
__copyright__ = "Copyright © 2020 Trent W. Buck"
__license__ = "expat"

__doc__ = """ build a Debian Live image that can install Debian 11 on ZFS 2

This uses mmdebstrap to do the heavy lifting;
it can run entirely without root privileges.
It emits a USB key disk image that contains a bootable EFI ESP,
which in turn includes a bootloader (refind), kernel, ramdisk, and filesystem.squashfs.

"""

parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
    "output_file", nargs="?", default=pathlib.Path("filesystem.img"), type=pathlib.Path
)
parser.add_argument(
    "--timezone",
    default="Australia/Melbourne",
    type=lambda s: s.split("/"),
    help='NOTE: MUST be "Area/Zone" not e.g. "UTC", for now',
)
parser.add_argument(
    "--locale", default="en_AU.UTF-8", help='NOTE: MUST end in ".UTF-8", for now'
)
args = parser.parse_args()


filesystem_img_size = "512M"  # big enough to include filesystem.squashfs + about 64M of bootloader, kernel, and ramdisk.
esp_offset = 1024 * 1024  # 1MiB
esp_label = "UEFI-ESP"  # max 8 bytes for FAT32
live_media_path = "debian-live"

with tempfile.TemporaryDirectory(prefix="debian-sid-zfs.") as td:
    td = pathlib.Path(td)
    subprocess.check_call(
        [
            "mmdebstrap",
            "--mode=unshare",
            "--variant=apt",
            '--aptopt=Acquire::http::Proxy "http://apt-cacher-ng.cyber.com.au:3142"',
            '--aptopt=Acquire::https::Proxy "DIRECT"',
            "--dpkgopt=force-unsafe-io",
            "--components=main contrib non-free",  # needed for CPU security patches
            "--include=init initramfs-tools xz-utils live-boot netbase",
            "--include=dbus",  # https://bugs.debian.org/814758
            "--include=linux-image-amd64 firmware-linux",
            # Have ZFS 2.0 support.
            "--include=zfs-dkms zfsutils-linux zfs-zed build-essential linux-headers-amd64",  # ZFS 2 support
            # Make the initrd a little smaller (41MB -> 20MB), at the expensive of significantly slower image build time.
            "--include=zstd",
            "--essential-hook=mkdir -p $1/etc/initramfs-tools/conf.d",
            "--essential-hook=>$1/etc/initramfs-tools/conf.d/zstd echo COMPRESS=zstd",
            # Be the equivalent of Debian Live GNOME
            # '--include=live-task-gnome',
            #'--include=live-task-xfce',
            # FIXME: enable this?  It makes live-task-xfce go from 1G to 16G... so no.
            #'--aptopt=Apt::Install-Recommends "true"',
            # ...cherry-pick instead
            # UPDATE: debian-installer-launcher DOES NOT WORK because we don't load crap SPECIFICALLY into /live/installer, in the ESP.
            # UPDATE: network-manager-gnome DOES NOT WORK, nor is systemd-networkd auto-started... WTF?
            #         end result is no networking.
            #'--include=live-config user-setup sudo firmware-linux haveged',
            #'--include=calamares-settings-debian udisks2',  # 300MB weirdo Qt GUI debian installer
            #'--include=xfce4-terminal',
            # x86_64 CPUs are undocumented proprietary RISC chips that EMULATE a documented x86_64 CISC ISA.
            # The emulator is called "microcode", and is full of security vulnerabilities.
            # Make sure security patches for microcode for *ALL* CPUs are included.
            # By default, it tries to auto-detect the running CPU, so only patches the CPU of the build server.
            "--include=intel-microcode amd64-microcode iucode-tool",
            "--essential-hook=>$1/etc/default/intel-microcode echo IUCODE_TOOL_INITRAMFS=yes IUCODE_TOOL_SCANCPUS=no",
            "--essential-hook=>$1/etc/default/amd64-microcode echo AMD64UCODE_INITRAMFS=yes",
            "--dpkgopt=force-confold",  # Work around https://bugs.debian.org/981004
            # DHCP/DNS/SNTP clients...
            # FIXME: use live-config ?
            "--include=libnss-resolve libnss-myhostname systemd-timesyncd",
            "--customize-hook=chroot $1 cp -alf /lib/systemd/resolv.conf /etc/resolv.conf",  # This probably needs to happen LAST
            # FIXME: fix resolv.conf to point to resolved, not "copy from the build-time OS"
            # FIXME: fix hostname & hosts to not exist, not "copy from the build-time OS"
            "--customize-hook=systemctl --root=$1 enable systemd-networkd systemd-timesyncd",  # is this needed?
            # Run a DHCP client on *ALL* ifaces.
            # Consider network "up" (start sshd and local login prompt) when *ANY* (not ALL) ifaces are up.
            "--customize-hook=>$1/etc/systemd/network/up.network printf '%s\n'  '[Match]' Name='en*' '[Network]' DHCP=yes",  # try DHCP on all ethernet ifaces
            "--customize-hook=mkdir $1/etc/systemd/system/systemd-networkd-wait-online.service.d",
            "--customize-hook=>$1/etc/systemd/system/systemd-networkd-wait-online.service.d/any-not-all.conf printf '%s\n' '[Service]' 'ExecStart=' 'ExecStart=/lib/systemd/systemd-networkd-wait-online --any'",
            # Hope there's a central smarthost SMTP server called "mail" in the local search domain.
            # FIXME: can live-config do this?
            "--include=msmtp-mta",
            "--customize-hook=>$1/etc/msmtprc printf '%s\n' 'account default' 'syslog LOG_MAIL' 'host mail' 'auto_from on'",
            # Hope there's a central RELP logserver called "logserv" in the local domain.
            # FIXME: can live-config do this?
            "--include=rsyslog-relp",
            """--customize-hook=>$1/etc/rsyslog.conf printf '%s\n' 'module(load="imuxsock")' 'module(load="imklog")' 'module(load="omrelp")' 'action(type="omrelp" target="logserv" port="2514" template="RSYSLOG_SyslogProtocol23Format")'""",
            # Run self-tests on all discoverable hard disks, and (try to) email if something goes wrong.
            "--include=smartmontools bsd-mailx",
            "--customize-hook=>$1/etc/smartd.conf echo 'DEVICESCAN -n standby,15 -a -o on -S on -s (S/../../7/00|L/../01/./01) -t -H -m root -M once'",
            # For rarely-updated, rarely-rebooted SOEs, apply what security updates we can into transient tmpfs COW.
            # This CANNOT apply kernel security updates (though it will download them).
            # This CANNOT make the upgrades persistent across reboots (they re-download each boot).
            # FIXME: Would it be cleaner to set Environment=NEEDRESTART_MODE=a in
            #        apt-daily-upgrade.service and/or
            #        unattended-upgrades.service, so
            #        needrestart is noninteractive only when apt is noninteractive?
            "--include=unattended-upgrades needrestart",
            "--customize-hook=echo 'unattended-upgrades unattended-upgrades/enable_auto_updates boolean true' | chroot $1 debconf-set-selections",
            """--customize-hook=>$1/etc/needrestart/conf.d/unattended-needrestart.conf  echo '$nrconf{restart} = "a";'""",  # https://bugs.debian.org/894444
            # Do an apt update & apt upgrade at boot time (as well as @daily).
            # The lack of /etc/machine-id causes these to be implicitly enabled.
            # FIXME: use dropin in /etc.
            "--customize-hook=>>$1/lib/systemd/system/apt-daily.service		printf '%s\n' '[Install]' 'WantedBy=multi-user.target'",
            "--customize-hook=>>$1/lib/systemd/system/apt-daily-upgrade.service	printf '%s\n' '[Install]' 'WantedBy=multi-user.target'",
            # FIXME: add support for this stuff (for the non-live final install this happens via ansible):
            #
            #            unattended-upgrades
            #            smartd
            #            networkd (boot off ANY NIC, not EVERY NIC -- https://github.com/systemd/systemd/issues/9714)
            #            refind (bootloader config)
            #            misc safety nets
            #            double-check that mmdebstrap's machine-id support works properly
            # Bare minimum to let me SSH in.
            # FIXME: make this configurable.
            # FIXME: trust a CA certificate instead -- see Zero Trust SSH, Jeremy Stott, LCA 2020 <https://youtu.be/lYzklWPTbsQ>
            # WARNING: tinysshd does not support RSA, nor MaxStartups, nor sftp (unless you also install openssh-client, which is huge).
            # FIXME: double-check no host keys are baked into the image (openssh-server and dropbear do this).
            "--include=tinysshd rsync",
            "--essential-hook=install -dm700 $1/root/.ssh",
            '--essential-hook=echo "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIIapAZ0E0353DaY6xBnasvu/DOvdWdKQ6RQURwq4l6Wu twb@cyber.com.au (Trent W. Buck)" >$1/root/.ssh/authorized_keys',
            # Bare minimum to let me log in locally.
            # DO NOT use this on production builds!
            "--essential-hook=chroot $1 passwd --delete root",
            # Configure language (not needed to boot).
            # Racism saves a **LOT** of space -- something like 2GB for Debian Live images.
            # FIXME: use live-config instead?
            "--include=locales localepurge",
            f"--essential-hook=echo locales locales/default_environment_locale   select {args.locale}       | chroot $1 debconf-set-selections",
            f"--essential-hook=echo locales locales/locales_to_be_generated multiselect {args.locale} UTF-8 | chroot $1 debconf-set-selections",
            # FIXME: https://bugs.debian.org/603700
            "--customize-hook=chroot $1 sed -i /etc/locale.nopurge -e 's/^USE_DPKG/#ARGH#&/'",
            "--customize-hook=chroot $1 localepurge",
            "--customize-hook=chroot $1 sed -i /etc/locale.nopurge -e 's/^#ARGH#//'",
            # Removing documentation also saves a LOT of space.
            "--dpkgopt=path-exclude=/usr/share/doc/*",
            "--dpkgopt=path-exclude=/usr/share/info/*",
            "--dpkgopt=path-exclude=/usr/share/man/*",
            "--dpkgopt=path-exclude=/usr/share/omf/*",
            "--dpkgopt=path-exclude=/usr/share/help/*",
            "--dpkgopt=path-exclude=/usr/share/gnome/help/*",
            # Configure timezone (not needed to boot)`
            # FIXME: use live-config instead?
            "--include=tzdata",
            f"--essential-hook=echo tzdata tzdata/Areas                    select {args.timezone[0]} | chroot $1 debconf-set-selections",
            f"--essential-hook=echo tzdata tzdata/Zones/{args.timezone[0]} select {args.timezone[1]} | chroot $1 debconf-set-selections",
            # Do the **BARE MINIMUM** to make a USB key that can boot on X86_64 UEFI.
            # We use mtools so we do not ever need root privileges.
            # We can't use mkfs.vfat, as that needs kpartx or losetup (i.e. root).
            # We can't use mkfs.udf, as that needs mount (i.e. root).
            # We can't use "refind-install --usedefault" as that runs mount(8) (i.e. root).
            # We don't use genisoimage because
            # 1) ISO9660 must die;
            # 2) incomplete UDF 1.5+ support;
            # 3) resulting filesystem can't be tweaked after flashing (e.g. debian-live/site.dir/etc/systemd/network/up.network).
            #
            # We use refind because 1) I hate grub; and 2) I like refind.
            # If you want aarch64 or ia32 you need to install their BOOTxxx.EFI files.
            # If you want kernel+initrd on something other than FAT, you need refind/drivers_xxx/xxx_xxx.EFI.
            #
            # FIXME: with qemu in UEFI mode (OVMF), I get dumped into startup.nsh (UEFI REPL).
            #        From there, I can manually type in "FS0:\EFI\BOOT\BOOTX64.EFI" to start refind, tho.
            #        So WTF is its problem?  Does it not support fallback bootloader?
            "--include=refind parted mtools",
            "--essential-hook=echo refind refind/install_to_esp boolean false | chroot $1 debconf-set-selections",
            "--customize-hook=echo refind refind/install_to_esp boolean true  | chroot $1 debconf-set-selections",
            "--customize-hook=chroot $1 mkdir -p /boot/USB /boot/EFI/BOOT",
            "--customize-hook=chroot $1 cp /usr/share/refind/refind/refind_x64.efi /boot/EFI/BOOT/BOOTX64.EFI",
            "--customize-hook=chroot $1 cp /usr/share/refind/refind/refind.conf-sample /boot/EFI/BOOT/refind.conf",
            f"--customize-hook=chroot $1 truncate --size={filesystem_img_size} /boot/USB/filesystem.img",
            f"--customize-hook=chroot $1 parted --script --align=optimal /boot/USB/filesystem.img  mklabel gpt  mkpart {esp_label} {esp_offset}b 100%  set 1 esp on",
            f"--customize-hook=chroot $1 mformat -i /boot/USB/filesystem.img@@{esp_offset} -F -v {esp_label}",
            f"--customize-hook=chroot $1 mmd     -i /boot/USB/filesystem.img@@{esp_offset} ::{live_media_path}",
            f"""--customize-hook=echo '"Boot with default options" "boot=live live-media-path={live_media_path}"' >$1/boot/refind_linux.conf""",
            f"""--customize-hook=chroot $1 find /boot/ -xdev -mindepth 1 -maxdepth 1 -not -name filesystem.img -not -name USB -exec mcopy -vsbpm -i /boot/USB/filesystem.img@@{esp_offset} {{}} :: ';'""",
            # FIXME: copy-out doesn't handle sparseness, so is REALLY slow (about 50 seconds).
            # Therefore instead leave it in the squashfs, and extract it later.
            #  f'--customize-hook=copy-out /boot/USB/filesystem.img /tmp/',
            #  f'--customize-hook=chroot $1 rm /boot/USB/filesystem.img',
            "sid",
            td / "filesystem.squashfs",
        ]
    )

    with args.output_file.open("wb") as f:
        subprocess.check_call(
            ["rdsquashfs", "--cat=boot/USB/filesystem.img", td / "filesystem.squashfs"],
            stdout=f,
        )
    subprocess.check_call(
        [
            "mcopy",
            "-i",
            f"{args.output_file}@@{esp_offset}",
            td / "filesystem.squashfs",
            f"::{live_media_path}/filesystem.squashfs",
        ]
    )