#!/usr/bin/env python3

import os
import shutil
import subprocess
import sys

IMAGES_DIRECTORY = '/images'
ADDONS_DIRECTORY = f"{IMAGES_DIRECTORY}/boot/addons"

ADDON_PATH                     = f"{ADDONS_DIRECTORY}/extra.commandline.addon.efi"
DEVICE_SECURE_BOOT_CERTIFICATE = f"{ADDONS_DIRECTORY}/secure-boot.pem"
DEVICE_SECURE_BOOT_PRIVATE_KEY = f"{ADDONS_DIRECTORY}/secure-boot.priv"
KERNEL_COMMANDLINE_FILE        = f"{ADDONS_DIRECTORY}/commandline"

PUAVO_INSTALL_LIBDIR = '/usr/lib/puavo-ltsp-install'


def log_info(message: str):
  print(message)


def log_error(message: str):
  print(f"error: {message}", file=sys.stderr)


def run_command(*args, check: bool = True, capture: bool = True):
  return subprocess.run(
    args,
    check=check,
    text=True,
    capture_output=capture,
  )


def replace_if_changed(destination: str, source: str):
  try:
    import filecmp
    if (
      os.path.exists(destination) and
        filecmp.cmp(destination, source, shallow=False)
    ):
      os.remove(source)
      return
  except Exception:
    # Fall through to moving the file
    pass

  shutil.move(source, destination)


def build_commandline_addon(commandline: str):
  os.makedirs(ADDONS_DIRECTORY, exist_ok=True)

  if not (os.path.isfile(DEVICE_SECURE_BOOT_CERTIFICATE) and \
          os.path.isfile(DEVICE_SECURE_BOOT_PRIVATE_KEY)):
    raise RuntimeError(
      f"device Secure Boot files are missing: {DEVICE_SECURE_BOOT_CERTIFICATE}, {DEVICE_SECURE_BOOT_PRIVATE_KEY}"
    )

  # Systemd-stub rejects an actually empty section, use a single space
  # in that case
  section_value = commandline if commandline.strip() != '' else ' '

  temporary_commandline_addon = f"{ADDON_PATH}.tmp"

  # Build UKI addon containing only a .cmdline section
  try:
    run_command(
      'ukify', 'build',
      f'--secureboot-private-key={DEVICE_SECURE_BOOT_PRIVATE_KEY}',
      f'--secureboot-certificate={DEVICE_SECURE_BOOT_CERTIFICATE}',
      f'--section=.cmdline:{section_value}',
      f'--output={temporary_commandline_addon}'
    )
  except subprocess.CalledProcessError as error:
    log_error(f"failed to build command-line addon (exit code {error.returncode}): {error.stderr.strip()}")
    raise RuntimeError("failed to build command-line addon")

  replace_if_changed(ADDON_PATH, temporary_commandline_addon)


def update_kernel_commandline(expected_kernel_commandline: str) -> int:
  current_kernel_commandline = None
  try:
    with open(KERNEL_COMMANDLINE_FILE, 'r') as file:
      current_kernel_commandline = file.read().strip()
  except FileNotFoundError:
    pass

  if expected_kernel_commandline == current_kernel_commandline:
    log_info("kernel command-line unchanged, nothing to do")
    return 0

  log_info("kernel command-line changed, rebuilding addon")
  log_info(f"old: {current_kernel_commandline}")
  log_info(f"new: {expected_kernel_commandline}")

  build_commandline_addon(expected_kernel_commandline)

  # Install the new command-line addon
  try:
    run_command(
      'puavo-install-and-update-uki-images',
      '--images-dir', IMAGES_DIRECTORY
    )
  except subprocess.CalledProcessError as error:
    log_error(error.stderr.strip() or "failed to install command line addon")
    return 1

  # Update the command line file with the new commandline
  try:
    with open(KERNEL_COMMANDLINE_FILE, 'w') as file:
      file.write(expected_kernel_commandline)
  except IOError as error:
    log_error(f"failed to write kernel command-line file: {error}")
    return 1

  log_info("kernel command-line addon updated")
  return 0


def clean_temporary_files():
  for path in [f"{ADDON_PATH}.tmp", f"{KERNEL_COMMANDLINE_FILE}.tmp"]:
    try:
      os.remove(path)
    except FileNotFoundError:
      pass


def main(new_kernel_commandline: str) -> int:
  try:
    return update_kernel_commandline(new_kernel_commandline.strip())
  except Exception as error:
    log_error(str(error))
    clean_temporary_files()
    return 1


if len(sys.argv) != 2:
  sys.exit(1)

new_kernel_commandline = sys.argv[1]

if __name__ == '__main__':
  sys.exit(main(new_kernel_commandline))
