#!/usr/bin/env python3

#

# This script inspects a given json proving a list of addons, and

# creates an addon for each key/value pair matching the given uki, distro and

# arch provided in input.

#

# Usage: python uki_create_addons.py input_json out_dir uki distro arch

#

# This tool requires the systemd-ukify and systemd-boot packages.

#

# Addon file

#-----------

# Each addon terminates with .addon

# Each addon contains only two types of lines:

# Lines beginning with '#' are description and thus ignored

# All other lines are command line to be added.

# The name of the end resulting addon is taken from the json hierarchy.

# For example, and addon in json['virt']['rhel']['x86_64']['hello.addon'] will

# result in an UKI addon file generated in out_dir called

# hello-virt.rhel.x86_64.addon.efi

#

# The common key, present in any sub-dict in the provided json (except the leaf dict)

# is used as place for default addons when the same addon is not defined deep

# in the hierarchy. For example, if we define test.addon (text: 'test1\n') in

# json['common']['test.addon'] = ['test1\n'] and another test.addon (text: test2) in

# json['virt']['common']['test.addon'] = ['test2'], any other uki except virt

# will have a test.addon.efi with text "test1", and virt will have a

# test.addon.efi with "test2"

#

# sbat.conf

#----------

# This dict is containing the sbat string for *all* addons being created.

# This dict is optional, but when used has to be put in a sub-dict with

# { 'sbat' : { 'sbat.conf' : ['your text here'] }}

# It follows the same syntax as the addon files, meaning '#' is comment and

# the rest is taken as sbat string and feed to ukify.

import os

import sys

import json

import collections

import subprocess

UKIFY_PATH = '/usr/lib/systemd/ukify'

def usage(err):

    print(f'Usage: {os.path.basename(__file__)} input_json output_dir uki distro arch')

    print(f'Error:{err}')

    sys.exit(1)

def check_clean_arguments(input_json, out_dir):

    # Remove end '/'

    if out_dir[-1:] == '/':

        out_dir = out_dir[:-1]

    if not os.path.isfile(input_json):

        usage(f'input_json {input_json} is not a file, or does not exist!')

    if not os.path.isdir(out_dir):

        usage(f'out_dir_dir {out_dir} is not a dir, or does not exist!')

    return out_dir

UKICmdlineAddon = collections.namedtuple('UKICmdlineAddon', ['name', 'cmdline'])

uki_addons_list = []

uki_addons = {}

addon_sbat_string = None

def parse_lines(lines, rstrip=True):

    cmdline = ''

    for l in lines:

        l = l.lstrip()

        if not l:

            continue

        if l[0] == '#':

            continue

        # rstrip is used only for addons cmdline, not sbat.conf, as it replaces

        # return lines with spaces.

        if rstrip:

            l = l.rstrip() + ' '

        cmdline += l

    if cmdline == '':

        return ''

    return cmdline

def parse_all_addons(in_obj):

    global addon_sbat_string

    for el in in_obj.keys():

        # addon found: copy it in our global dict uki_addons

        if el.endswith('.addon'):

            uki_addons[el] = in_obj[el]

    if 'sbat' in in_obj and 'sbat.conf' in in_obj['sbat']:

        # sbat.conf found: override sbat with the most specific one found

        addon_sbat_string = parse_lines(in_obj['sbat']['sbat.conf'], rstrip=False)

def recursively_find_addons(in_obj, folder_list):

    # end of recursion, leaf directory. Search all addons here

    if len(folder_list) == 0:

        parse_all_addons(in_obj)

        return

    # first, check for common folder

    if 'common' in in_obj:

        parse_all_addons(in_obj['common'])

    # second, check if there is a match with the searched folder

    if folder_list[0] in in_obj:

        folder_next = in_obj[folder_list[0]]

        folder_list = folder_list[1:]

        recursively_find_addons(folder_next, folder_list)

def parse_in_json(in_json, uki_name, distro, arch):

    with open(in_json, 'r') as f:

        in_obj = json.load(f)

    recursively_find_addons(in_obj, [uki_name, distro, arch])

    for addon_name, cmdline in uki_addons.items():

        addon_name = addon_name.replace(".addon","")

        addon_full_name = f'{addon_name}-{uki_name}.{distro}.{arch}.addon.efi'

        cmdline = parse_lines(cmdline).rstrip()

        if cmdline:

            uki_addons_list.append(UKICmdlineAddon(addon_full_name, cmdline))

def create_addons(out_dir):

    for uki_addon in uki_addons_list:

        out_path = os.path.join(out_dir, uki_addon.name)

        cmd = [

            f'{UKIFY_PATH}', 'build',

            f'--cmdline="{uki_addon.cmdline}"',

            f'--output={out_path}']

        if addon_sbat_string:

            cmd.append('--sbat="' + addon_sbat_string.rstrip() +'"')

        subprocess.check_call(cmd, text=True)

if __name__ == "__main__":

    argc = len(sys.argv) - 1

    if argc != 5:

        usage('too few or too many parameters!')

    input_json = sys.argv[1]

    out_dir = sys.argv[2]

    uki_name = sys.argv[3]

    distro = sys.argv[4]

    arch = sys.argv[5]

    out_dir = check_clean_arguments(input_json, out_dir)

    parse_in_json(input_json, uki_name, distro, arch)

    create_addons(out_dir)