sysmodule_test.py

#!/usr/bin/python3

import os, sys
import pyudev
import subprocess
import time
import shutil
import argparse
from pathlib import Path

def get_fusee_path(*p):
    return os.path.join(os.environ["FUSEE"], *p)

def run_command(prog, args):
    return subprocess.Popen([prog, *args], stdout=subprocess.PIPE)

def run_payload(path):
    return run_command("python", [get_fusee_path("fusee-launcher.py"), "-w", path]).wait()

def boot_switch():
    return run_payload(get_fusee_path("fusee-primary.bin"))

def run_memloader():
    return run_payload(get_fusee_path("memloader.bin"))

def find_device():
    time.sleep(5)
    ctx = pyudev.Context()
    removable = [
        device for device in ctx.list_devices(subsystem='block', DEVTYPE='disk')
        if device.attributes.asstring('removable') == "1"
    ][0]
    return [
        device.device_node
        for device in ctx.list_devices(subsystem='block', DEVTYPE='partition', parent=removable)
    ][0]

def find_mountpoint(device):
    return run_command("lsblk", ["-n", "-oMOUNTPOINT", device]).communicate()[0].decode().strip("\n")

def mount_switch(device):
    return run_command("udisksctl", ["mount", "-b", device]).wait()

def unmount_switch(device):
    return run_command("udisksctl", ["unmount", "-b", device]).wait()

def main(argc, argv):
    parser = argparse.ArgumentParser(description="Install a sysmodule and reboot the console.")
    parser.add_argument("-p", "--path", required=True, dest="path", help="Path of the executable to be copied.")
    parser.add_argument("-t", "--title_id", required=True, dest="tid")

    if argc < 2:
        parser.print_help()
        return 1

    args = parser.parse_args(argv[1:])
    path = Path(args.path)

    print("Running memloader...")
    rc = run_memloader()
    if rc:
        print(f"Error: {rc:#x}")
        return rc

    device = find_device()
    if device == "":
        print("Couldn't find device")
        return 1

    print("Mounting console...")
    rc = mount_switch(device)
    if rc:
        print(f"Error: {rc:#x}")
        return rc

    print("Finding mountpoint...")
    mountpoint = find_mountpoint(device)
    if mountpoint == "":
        print("Couldn't find mount point")
        return 1

    args.tid = int(args.tid.strip("0x"), 16)
    dest = Path(mountpoint, "atmosphere", "contents", f"{args.tid:016X}", "exefs.nsp")
    print(f"Copying from {path} to {dest}...")
    Path(dest.parent, "flags").mkdir(parents=True, exist_ok=True)
    shutil.copy2(path, dest)
    Path(dest.parent, "flags", "boot2.flag").touch()
    time.sleep(1)

    print("Unmounting console...")
    rc = unmount_switch(device)
    if rc:
        print(f"Error: {rc:#x}")
        return rc

    print("You have 10 seconds to reboot the console.")
    time.sleep(10)

    return boot_switch()

if __name__ == "__main__":
    sys.exit(main(len(sys.argv), sys.argv))