#! /usr/bin/python

import subprocess
import os

#fail if not root
if subprocess.check_output("id -u", shell=True).decode('utf-8').strip() != '0':
    print("You need to be root to run this script")
    exit(1)
variables = {}
# import /boot/env.txt and parse it
print("Generating extlinux config")
try:
    with open('/boot/env.txt', 'r') as f:
        config = map(lambda x: x.strip(), f.readlines())
        config = filter(lambda x: not x.startswith('#'), config)
        config = list(filter(lambda x: not x == '', config))
        config = [x.replace('"','') for x in config]
        config = [x.replace("'",'') for x in config]
        print(config)
        for line in config:
            key, value = line.split('=', 1)
            variables[key] = value
except FileNotFoundError:
    print("No /boot/env.txt found, install u-boot/towboot for your device or create /boot/env.txt manually")
    print("See /usr/share/doc/update-extlinux/examples/env.txt for an example")
try:
    if not 'FDT' in variables:
        extlinux_conf = f'''# Generated by update-extlinux 
        # Do not edit this file directly, edit /boot/env.txt instead and run /usr/bin/update-extlinux
        LABEL {variables['LABEL']}
            KERNEL {variables['KERNEL']}
            INITRD {variables['INITRD']}
        '''
    else:
        extlinux_conf = f'''# Generated by update-extlinux 
# Do not edit this file directly, edit /boot/env.txt instead and run /usr/bin/update-extlinux
LABEL {variables['LABEL']}
    KERNEL {variables['KERNEL']}
    INITRD {variables['INITRD']}
    FDT {variables['FDT']}
'''
except KeyError:
    print("You need to set LABEL, KERNEL and INITRD in /boot/env.txt")
    exit(1)
# fdtoverlays is optional append it if it's set
try:
    extlinux_conf += f"    FDTOVERLAYS {variables['FDTOVERLAYS']}\n"
except KeyError:
    pass

# if old extlinux.conf exists, make a backup
if os.path.exists("/boot/extlinux/extlinux.conf"):
    print("Making backup of /boot/extlinux/extlinux.conf")
    print("You can restore it with cp /boot/extlinux/extlinux.conf.bak /boot/extlinux/extlinux.conf")
    subprocess.run("cp /boot/extlinux/extlinux.conf /boot/extlinux/extlinux.conf.bak", shell=True)
# if /boot/extlinux doesn't exist, create it
if not os.path.exists("/boot/extlinux"):
    print("Creating /boot/extlinux")
    os.mkdir("/boot/extlinux")
with open("/boot/extlinux/extlinux.conf", 'w') as f:
    # if variables['ROOT'] is not set, use current root
    if not 'ROOT' in variables:
        dev = subprocess.check_output("df -h / | tail -n 1 | awk '{print $1}'", shell=True).decode('utf-8').strip()
    else:
        dev = variables['ROOT']
    uuid=subprocess.check_output(f"blkid -s UUID -o value {dev}", shell=True).decode('utf-8').strip()
    fs = subprocess.check_output(f"blkid -s TYPE -o value {dev}", shell=True).decode('utf-8').strip()
    f.write(extlinux_conf)
    # if dev has /dev/mapper/luks- in it, then it's encrypted
    if dev.startswith('/dev/mapper/luks-'):
        print("Root {dev} is encrypted")
        real_dev=subprocess.check_output(f"lsblk -sJp | jq -r --arg dsk '{dev}' '.blockdevices | .[] | select(.name == $dsk) | .children | .[0] | .name'",shell=True).decode('utf-8').strip()
        real_uuid=subprocess.check_output(f"blkid -s UUID -o value {real_dev}", shell=True).decode('utf-8').strip()
        if fs == 'btrfs':
            f.write(f"    APPEND root=UUID={uuid} rootflags=subvol=@ rw cryptdevice=UUID={real_uuid}:luks-{real_uuid} root={dev} rootwait {variables['CMDLINE']}")
        else:
            f.write(f"    APPEND root=UUID={uuid} rw cryptdevice=UUID={real_uuid}:luks-{real_uuid} root={dev}rootwait {variables['CMDLINE']}")
    else: 
        if fs == 'btrfs':
            f.write(f"    APPEND root=UUID={uuid} rootflags=subvol=@ rw rootwait {variables['CMDLINE']}")
        else:
            f.write(f"    APPEND root=UUID={uuid} rw rootwait {variables['CMDLINE']}")
print("Done")
