post_install_securebootify.py

import os
import subprocess
from typing import Tuple
from pathlib import Path
import sys
from getpass import getpass
from shutil import rmtree

def is_root() -> bool:
    """
    Check if the script is running as root
    
    :return: True if root, False otherwise
    :rtype: bool
    """
    return os.geteuid() == 0 or os.getuid() == 0

def is_arch_based() -> bool:
    """
    Check if the OS is arch based
    
    :return: True if likely Arch based, False otherwise
    :rtype: bool
    """

    try:
        subprocess.run(["pacman", "--version"])
    except Exception:
        return False

    return True

def create_mok_keys() -> Tuple[str]:
    """
    Generate mok keys
    
    :return: Path to generated MOK.cer, MOK.crt and MOK.key (in that order)
    :rtype: Tuple[str]
    """
    
    cmd1 = 'openssl req -x509 -newkey rsa:2048 -keyout MOK.key -out MOK.crt -subj "/CN=Aero/"' # Replace Aero with your name!
    cmd2 = "openssl x509 -in MOK.crt -out MOK.cer -outform DER"

    ret1 = subprocess.run(cmd1, shell=True, capture_output=True)
    if ret1.returncode != 0:
        print(f"Error: {ret1.stderr.decode('utf-8')}")
        exit(1)
    
    ret2 = subprocess.run(cmd2, shell=True, capture_output=True)
    if ret2.returncode != 0:
        print(f"Error: {ret2.stderr.decode('utf-8')}")
        exit(1)

    return Path(os.getcwd()).joinpath("MOK.cer").absolute(), Path(os.getcwd()).joinpath("MOK.crt").absolute(), Path(os.getcwd()).joinpath("MOK.key").absolute()

def sign_bootmanager(keypath: str, crtpath: str):
    """
    Sign the bootloader
    
    :param keypath: Path to .key file
    :type keypath: str
    :param crtpath: Path to .crt file
    :type crtpath: str
    """

    if (not Path(keypath).exists()) or (not Path(crtpath).exists()):
        print("Provided paths to certificate files for signing do not exist. Exiting...")
        sys.exit(1)

    ret2 = subprocess.run(f'sbsign --key {keypath} --cert {crtpath} --output /boot/efi/EFI/BOOT/BOOTx64.EFI /boot/efi/EFI/BOOT/BOOTx64.EFI', shell=True, capture_output=True)
    if ret2.returncode != 0:
        print(f"Error: {ret2.stderr.decode('utf-8')}")
        exit(1)

def initcpico_setup(keypath: str, crtpath: str):
    path = "/etc/initcpio/post/kernel-sbsign"
    cmd = "chmod +x /etc/initcpio/post/kernel-sbsign"
    
    if (not Path(keypath).exists()) or (not Path(crtpath).exists()):
        print("Provided paths to certificate files for signing do not exist. Exiting...")
        sys.exit(1)

    content = [
        "#!/usr/bin/env bash\n",
        'kernel="$1"\n',
        '[[ -n "$kernel" ]] || exit 0\n',

        "# use already installed kernel if it exists\n",
        '[[ ! -f "$KERNELDESTINATION" ]] || kernel="$KERNELDESTINATION"\n\n',

        f'keypairs=({keypath} {crtpath})\n'

        'for (( i=0; i<${#keypairs[@]}; i+=2 )); do\n',
        '    key="${keypairs[$i]}" cert="${keypairs[(( i + 1 ))]}"\n',
        '    if ! sbverify --cert "$cert" "$kernel" &>/dev/null; then\n',
        '        sbsign --key "$key" --cert "$cert" --output "$kernel" "$kernel"\n',
        '    fi\n',
        'done\n'
    ]

    with open(path, 'w') as f:
        f.writelines(content)

    subprocess.run(cmd)

if __name__ == '__main__':
    banner = """
    _                    _       ____                
   / \     _ __    ___  | |__   / ___|    ___    ___ 
  / _ \   | '__|  / __| | '_ \  \___ \   / _ \  / __|
 / ___ \  | |    | (__  | | | |  ___) | |  __/ | (__ 
/_/   \_\ |_|     \___| |_| |_| |____/   \___|  \___|


Secure Boot enabler for Arch and Arch based distros.
Script type: Repacker (Enable secure boot for ISO file)
"""

    if not is_arch_based():
        print("[ERROR]: This distro is not Arch or Arch based. Exiting...")
        sys.exit(1)
    
    print(banner)

    if not is_root():
        subprocess.run(["pkexec", sys.executable, __file__])
        sys.exit()
    
    cwd = os.getcwd()
    p = Path(cwd).joinpath("archsec-post-install")

    if p.exists():
        print(f"{p.absolute()} exisits. Exiting...")
        sys.exit(1)
    
    os.chdir(p.absolute())
    cer, crt, key = create_mok_keys()

    cmd = f"mokutil --import {cer}"
    print("[INFO]: A password is required to enroll the secure boot key.\nWhile the key itself won't be enrolled until next reboot (you should use enroll key), a password should be set now.\nThe password will be used when you enroll the key.\n\n")
    password = getpass(prompt="[ACTION]: Password:\t", echo_char="*")
    password2 = getpass(prompt="[ACTION]: Enter password again:\t", echo_char="*")

    if password != password2:
        print("Passwords do not match. Exiting...")
        sys.exit(1)
    
    subprocess.run(cmd, shell=True, input=password.encode())
    print("[INFO]: Remember to use the same password to enroll the keys on reboot!")

    initcpico_setup(key, crt)
    sign_bootmanager(key, crt)

    os.chdir(cwd)
    rmtree(path=p.absolute())

    print("Done")