#!/usr/bin/python3
"""
Helper to build Raspberry Pi images.

The PiRogue image generation tool would ideally use daily auto-built images
provided by Debian (https://raspi.debian.net/) but the build process might be
stalled for several weeks or months, so we're building our own image(s) in
a weekly fashion. This makes sure we don't end up shipping PiRogue images
containing severely outdated packages (meaning a possibly significant amount of
missing security/stability fixes).

This script is a wrapper around image-specs's Makefile, and it also helps
maintain the required associated sudo rules.

Generate sudo configuration for the specified image(s):

    ./raspi-build --sudo raspi_4_bookworm

Build, compress, checksum, and rsync the specified images(s) to a web server
(see RSYNC_DEST):

    ./raspi-build --build raspi_4_bookworm
"""

import argparse
import hashlib
import logging
import os
import subprocess
import sys
from pathlib import Path

# Settings:
WORK_DIR = Path('~/image-specs').expanduser()
RSYNC_DEST = 'website:pirogue.apt.debamax.com/raspi-images/'


def sudo(images):
    """
    Generate the sudo snippet for all specified images.
    """
    print(f'# Generated by {Path(__file__).resolve()}')
    for image in images:
        # Keep in sync with the Makefile in image-specs:
        print(f'{login} ALL=NOPASSWD: /usr/bin/vmdb2 --verbose --rootfs-tarball={image}.tar.gz '
              f'--output={image}.img {image}.yaml --log {image}.log')
        # Keep in sync with root-owned products/byproducts (also noting the
        # escaped ':' character):
        print(f'{login} ALL=NOPASSWD: /usr/bin/chown {uid}\\:{gid} {WORK_DIR}/{image}.img')
        print(f'{login} ALL=NOPASSWD: /usr/bin/chown {uid}\\:{gid} {WORK_DIR}/{image}.tar.gz')


def run(title, *cmd):
    """
    Log and exec.
    """
    logging.info(title)
    subprocess.check_call(*cmd)


def build(images):
    """
    Build all specified images.
    """
    os.chdir(WORK_DIR)
    for image in images:
        logging.info('=== BEGIN: WORK WITH %s ===', image)
        # Generate the uncompressed images, then fix permissions on our own:
        run('Generate image',
            ['make', f'{image}.img'])
        run('Fix image owner',
            ['sudo', 'chown', f'{uid}:{gid}', f'{WORK_DIR}/{image}.img'])
        run('Fix tarball owner',
            ['sudo', 'chown', f'{uid}:{gid}', f'{WORK_DIR}/{image}.tar.gz'])

        # Generate compressed image ourselves to avoid timestamp-related fun if
        # we were using the Makefile:
        run('Compress image',
            ['xz', '-9', f'{image}.img'])

        # Compute checksum ourselves (same reason as above), without using
        # run(): sha256sum doesn't take an output argument so we'd need
        # shell=True.
        logging.info('Generate checksum')
        sha256 = hashlib.file_digest(Path(f'{image}.img.xz').open('rb'), 'sha256')
        Path(f'{image}.img.xz.sha256').write_text(f'{sha256.hexdigest()}  {image}.img.xz\n')

        # Sync only what we need:
        run('Sync compressed image and checksum',
            ['rsync', '-av', *Path('.').glob(f'{image}.img.xz*'), RSYNC_DEST])
        logging.info('=== END: WORK WITH %s ===', image)

    logging.info('=== BEGIN: CLEAN-UP ===')
    run('Cleaning git',
        ['git', 'clean', '-xdf'])
    logging.info('=== END: CLEAN-UP ===')


if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO,
                        datefmt='%H:%M:%S',
                        format='%(asctime)s %(message)s')
    parser = argparse.ArgumentParser()
    parser.add_argument('images', metavar='IMAGE', type=str, nargs='+',
                        help='an image to build (e.g. raspi_4_bookworm)')
    parser.add_argument('--sudo', action='store_true',
                        help='generate the sudo snippet required to build specified images')
    parser.add_argument('--build', action='store_true',
                        help='build the specified images')
    args = parser.parse_args()

    login = os.getlogin()
    uid = os.getuid()
    gid = os.getgid()

    if args.sudo:
        sudo(args.images)
        sys.exit(0)

    if args.build:
        build(args.images)
        sys.exit(0)