setup-audio

#!/usr/bin/env python3

import argparse
import json
import os
import sys
from functions import *

# parse arguments from the cli. Only for testing/advanced use.
def process_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("-b", dest="board_name", type=str, nargs=1, default=[""],
                        help="Override board name.")
    parser.add_argument("--enable-debug", action='store_const', const="Enabling", dest="debug",
                        help="Enable audio debugging.")
    parser.add_argument("--disable-debug", action='store_const', const="Disabling", dest="debug",
                        help="Disable audio debugging.")
    parser.add_argument("--force-avs-install", action="store_true", dest="force_avs_install", default=False,
                        help="DANGEROUS: Force enable AVS install. MIGHT CAUSE PERMANENT DAMAGE TO SPEAKERS!")
    return parser.parse_args()


def avs_audio(board, username):
    if args.debug:
        print_status(f"{args.debug} AVS debugging")
        if args.debug == "Enabling":
            cpfile("conf/avs/snd-avs-dbg.conf", "/etc/modprobe.d/snd-avs-dbg.conf")
        else:
            rmfile("/etc/modprobe.d/snd-avs-dbg.conf")
        print_status("Done, please reboot for changes to take effect.")
        exit()

    print_status("Installing AVS")
    # Only show the warning to devices with max98357a
    override_avs = False
    if path_exists("/sys/bus/acpi/devices/MX98357A:00"):
        if args.force_avs_install:
            print_error(
                "WARNING: Your device has max98357a and can cause permanent damage to your speakers if you set the volume too loud!")
            user_input = input('Type "I understand the risk of permanently damaging my speakers" in all caps to continue: ')
            while user_input != "I UNDERSTAND THE RISK OF PERMANENTLY DAMAGING MY SPEAKERS":
                user_input = input(
                    'Type "I understand the risk of permanently damaging my speakers" in all caps to continue: ')
            override_avs = True
        else:
            print_error(
                "WARNING: Your device has max98357a and can cause permanent damage to your speakers if you "
                "set the volume too loud! As a safety precaution devices with max98357a have speakers "
                "disabled until a fix is in place. Headphones and HDMI audio are safe from this.")
            print_question("If you want to disable this check, restart the script with --force-avs-install")

            user_input = input(
                'Type "I Understand my speakers will not work since my device has max98357a!" in all caps to continue: ')
            while user_input != "I UNDERSTAND MY SPEAKERS WILL NOT WORK SINCE MY DEVICE HAS MAX98357A!":
                user_input = input(
                    'Type "I Understand my speakers will not work since my device has max98357a!" in all caps to continue: ')
            override_avs = False

    # Copy ucms
    print_status("Installing UCM configuration")
    rmdir("/tmp/chromebook-ucm-conf")
    bash("git clone https://github.com/WeirdTreeThing/chromebook-ucm-conf /tmp/chromebook-ucm-conf")
    cpdir("/tmp/chromebook-ucm-conf/avs", "/usr/share/alsa/ucm2/conf.d/")

    # avs tplg is from https://github.com/thesofproject/avs-topology-xml
    print_status("Installing topology")
    cpdir("conf/avs/tplg", "/lib/firmware/intel/avs")
    print_status("Installing modprobe config")
    cpfile("conf/avs/snd-avs.conf", "/etc/modprobe.d/snd-avs.conf")

    # Install wireplumber config for dmic if wireplumber is installed on the system
    if path_exists("/usr/bin/wireplumber"):
        print_status("Forcing avs_dmic to use S16LE format")
        mkdir("/etc/wireplumber/main.lua.d/", create_parents=True)
        cpfile("conf/avs/51-avs-dmic.lua", "/etc/wireplumber/main.lua.d/51-avs-dmic.lua")

    # updated avs dsp firmware recently got merged upstream but is not packaged in any distro yet
    print_status("Installing AVS firmware")
    mkdir("/lib/firmware/intel/avs/skl")
    mkdir("/lib/firmware/intel/avs/apl")
    urlretrieve("https://git.kernel.org/pub/scm/linux/kernel/git/firmware/linux-firmware.git/plain/intel/avs/apl/"
                "dsp_basefw.bin", filename="/lib/firmware/intel/avs/apl/dsp_basefw.bin")
    urlretrieve("https://git.kernel.org/pub/scm/linux/kernel/git/firmware/linux-firmware.git/plain/intel/avs/skl/"
                "dsp_basefw.bin", filename="/lib/firmware/intel/avs/skl/dsp_basefw.bin")

    # Delete topology for max98357a to prevent it from working until there is a volume limiter.
    if not override_avs:
        rmfile("/lib/firmware/intel/avs/max98357a-tplg.bin")

    if not path_exists(f"/lib/modules/{bash('uname -r')}/kernel/sound/soc/intel/avs"):
        print_error("Looks like your kernel doesn't have the avs modules. Make sure you are on at least 6.0 with avs enabled")
        exit(0)


def apl_audio(board, username):
    print_status("Apollolake has two audio drivers:")
    print_status("SOF: Stable but doesn't work with headphones.")
    print_status("AVS: Unstable and can cause damage to speakers but supports all audio hardware.")
    print_error("NOTE: Speakers are disabled on AVS as a safety precaution. (use --force-avs-install to override)"
                "Your speakers will still work on SOF though.")

    while True:
        user_input = input("Which driver would you like to use? [sof/avs]: ")
        if user_input.lower() == "sof":
            print_status("Using sof")
            # Remove avs modprobe config if it exists
            rmfile("/etc/modprobe.d/snd-avs.conf")
            sof_audio("apl")
            # Install apl specific modprobe config
            cpfile("conf/sof/apl-sof.conf", "/etc/modprobe.d/apl-sof.conf")
            break
        elif user_input.lower() == "avs":
            print_status("Using avs")
            # Remove sof modprobe config if it exists
            rmfile("/etc/modprobe.d/snd-sof.conf")
            rmfile("/etc/modprobe.d/apl-sof.conf")
            avs_audio(board, username)
            break
        else:
            print_error(f"Invalid option: {user_input}")
            continue


def sof_audio(platform):
    if args.debug:
        print_status(f"{args.debug} SOF debugging")
        if args.debug == "Enabling":
            cpfile("conf/sof/alsa-sof-dbg.conf", "/etc/modprobe.d/snd-sof-dbg.conf")
        else:
            rmfile("/etc/modprobe.d/snd-sof-dbg.conf")
        print_status("Done, please reboot for changes to take effect.")
        exit()

    print_status("Installing SOF")

    # Install required packages
    print_status("Installing SOF firmware")
    install_package("sof-firmware", "firmware-sof-signed", "alsa-sof-firmware", "sof-firmware", "sof-firmware")
    # linux-firmware has different names on ubuntu (and ubuntu based like PopOS) and debian
    with open("/etc/os-release") as f:
        if "ubuntu" in f.read():
            install_package(deb_package="linux-firmware")
        else:
            install_package("linux-firmware", "firmware-linux-free firmware-linux-nonfree", "linux-firmware",
                            "linux-firmware", "linux-firmware")
    install_package("alsa-utils", "alsa-utils", "alsa-utils", "alsa-utils", "alsa-utils")

    # Force sof driver
    print_status("Installing modprobe config")
    cpfile("conf/sof/snd-sof.conf", "/etc/modprobe.d/snd-sof.conf")
    
    print_status("Installing UCM configuration")
    rmdir("/tmp/chromebook-ucm-conf")
    bash("git clone https://github.com/WeirdTreeThing/chromebook-ucm-conf /tmp/chromebook-ucm-conf")
    match platform:
        case "apl":
            cpdir("/tmp/chromebook-ucm-conf/apl/sof-bxtda7219ma", "/usr/share/alsa/ucm2/conf.d/sof-bxtda7219ma")
        case "glk":
            cpdir("/tmp/chromebook-ucm-conf/glk/sof-glkda7219ma", "/usr/share/alsa/ucm2/conf.d/sof-glkda7219ma")
            cpdir("/tmp/chromebook-ucm-conf/glk/sof-cs42l42", "/usr/share/alsa/ucm2/conf.d/sof-cs42l42")
            cpdir("/tmp/chromebook-ucm-conf/glk/sof-glkrt5682ma", "/usr/share/alsa/ucm2/conf.d/sof-glkrt5682ma")
        case "cml":
            cpdir("/tmp/chromebook-ucm-conf/cml/sof-rt5682", "/usr/share/alsa/ucm2/conf.d/sof-rt5682")
            cpdir("/tmp/chromebook-ucm-conf/cml/sof-cmlda7219ma", "/usr/share/alsa/ucm2/conf.d/sof-cmlda7219ma")
            cpdir("/tmp/chromebook-ucm-conf/cml/sof-cml_rt1011_", "/usr/share/alsa/ucm2/conf.d/sof-cml_rt1011_")
            cpdir("/tmp/chromebook-ucm-conf/cml/sof-cml_max9839", "/usr/share/alsa/ucm2/conf.d/sof-cml_max9839")
        case "tgl":
            cpdir("/tmp/chromebook-ucm-conf/tgl/sof-rt5682", "/usr/share/alsa/ucm2/conf.d/sof-rt5682")
        case "jsl":
            cpdir("/tmp/chromebook-ucm-conf/jsl/sof-rt5682", "/usr/share/alsa/ucm2/conf.d/sof-rt5682")
            cpdir("/tmp/chromebook-ucm-conf/jsl/sof-da7219max98", "/usr/share/alsa/ucm2/conf.d/sof-da7219max98")
            cpdir("/tmp/chromebook-ucm-conf/jsl/sof-cs42l42", "/usr/share/alsa/ucm2/conf.d/sof-cs42l42")
            # JSL needs tplg build from upstream which have not been shipped in distros yet
            cpdir("conf/sof/tplg", "/lib/firmware/intel/sof-tplg")
        case "adl":
            cpdir("/tmp/chromebook-ucm-conf/adl/sof-rt5682", "/usr/share/alsa/ucm2/conf.d/sof-rt5682")
            cpdir("/tmp/chromebook-ucm-conf/adl/sof-nau8825", "/usr/share/alsa/ucm2/conf.d/sof-nau8825")
            cpdir("/tmp/chromebook-ucm-conf/adl/sof-ssp_amp", "/usr/share/alsa/ucm2/conf.d/sof-ssp_amp")

    # Common dmic split ucm
    cpdir("/tmp/chromebook-ucm-conf/dmic-common", "/usr/share/alsa/ucm2/conf.d/dmic-common")

    # Common hdmi configurations
    cpdir("/tmp/chromebook-ucm-conf/hdmi-common", "/usr/share/alsa/ucm2/conf.d/hdmi-common")


def bsw_audio():
    if args.debug:
        print_status(f"{args.debug} SOF BSW debugging")
        if args.debug == "Enabling":
            cpfile(f"conf/sof/bsw-sof-dbg.conf", "/etc/modprobe.d/bsw-sof-dbg.conf")
        else:
            rmfile("/etc/modprobe.d/bsw-sof-dbg.conf")
        print_status("Done, please reboot for changes to take effect.")
        exit()

    print_status("Fixing braswell/baytrail/broadwell audio")
    install_package("alsa-ucm-conf", "alsa-ucm-conf", "alsa-ucm-conf", "alsa-ucm-conf", "alsa-ucm-conf")
    install_package("sof-firmware", "firmware-sof-signed", "alsa-sof-firmware", "sof-firmware", "sof-firmware")
    cpfile("conf/sof/bsw-sof.conf", "/etc/modprobe.d/bsw-sof.conf")


def amd_audio(amd_platform):
    # amd_platform is an int that represents the generation
    # 0 = stoneyridge (grunt)
    # 1 = picasso/zen+ (zork)
    # 2 = cezanne/zen3 (guybrush)
    # 3 = mendocino/zen2 (skyrim)
    print_status("Installing UCM configuration")
    rmdir("/tmp/chromebook-ucm-conf")
    bash("git clone https://github.com/WeirdTreeThing/chromebook-ucm-conf /tmp/chromebook-ucm-conf")
    match amd_platform:
        case 0:
            cpdir("/tmp/chromebook-ucm-conf/stoney/acpd7219m98357", "/usr/share/alsa/ucm2/conf.d/acpd7219m98357")
        case 1:
            cpdir("/tmp/chromebook-ucm-conf/picasso/acp3xalc5682m98", "/usr/share/alsa/ucm2/conf.d/acp3xalc5682m98")
            cpdir("/tmp/chromebook-ucm-conf/picasso/acp3xalc5682101", "/usr/share/alsa/ucm2/conf.d/acp3xalc5682101")
        case 2:
            cpdir("/tmp/chromebook-ucm-conf/cezanne/sof-rt5682s-rt1", "/usr/share/alsa/ucm2/conf.d/sof-rt5682s-rt1")
        case 3:
            cpdir("/tmp/chromebook-ucm-conf/mendocino/sof-rt5682s-hs-", "/usr/share/alsa/ucm2/conf.d/sof-rt5682s-hs-")

    print_status("After rebooting you will need to set speakers as the output device in sound settings.")


if __name__ == "__main__":
    if os.geteuid() == 0 and not path_exists("/tmp/username"):
        print_error("Please start the script as non-root/without sudo")
        exit(1)

    args = process_args()  # process args before elevating to root for better ux

    # Restart script as root
    if os.geteuid() != 0:
        # save username
        with open("/tmp/username", "w") as file:
            file.write(os.getlogin())

        # restart script as root
        # make the two people that use doas happy
        if path_exists("/usr/bin/doas"):
            doas_args = ['doas', sys.executable] + sys.argv + [os.environ]
            os.execlpe('doas', *doas_args)
        # other 99 percent of linux users
        sudo_args = ['sudo', sys.executable] + sys.argv + [os.environ]
        os.execlpe('sudo', *sudo_args)

    # read username
    with open("/tmp/username", "r") as file:
        user_id = file.read()

    # Important message
    print_warning("WARNING: this audio script is not fully functional yet!")

    if args.board_name[0] == "":
        # Get the board name from dmi
        with open("/sys/devices/virtual/dmi/id/product_name", "r") as dmi:
            device_board = dmi.read().lower().strip()
    else:  # use the board name from the args, for testing only
        device_board = str(args.board_name[0]).lower().strip()
        print_warning(f"Board name overriden to: {device_board}")

    with open("conf/board-generations.json", "r") as file:
        boards = json.load(file)

    try:
        match boards[device_board]:
            case "bdw":
                bsw_audio()
            case "byt" | "bsw":
                bsw_audio()
            case "skl" | "kbl":
                avs_audio(device_board, user_id)
            case "apl":
                apl_audio(device_board, user_id)
            case "glk" | "cml" | "jsl" | "tgl" | "adl":
                sof_audio(boards[device_board])
            case "stoney":
                amd_audio(0)
            case "picasso":
                amd_audio(1)
            case "cezanne":
                amd_audio(2)
            case "mendocino":
                amd_audio(3)
            case _:
                print_error(f"Unknown/Unsupported chromebook model: {device_board}")
                exit(1)
    except KeyError:
        print_error(f"Unknown/Unsupported chromebook model: {device_board}")
        exit(1)

    # Install wireplumber config to increase headroom
    # fixes instability and crashes on various devices
    if path_exists("/usr/bin/wireplumber"):
      print_status("Increasing alsa headroom (fixes instability)")
      mkdir("/etc/wireplumber/main.lua.d/", create_parents=True)
      cpfile("conf/common/51-increase-headroom.lua", "/etc/wireplumber/main.lua.d/51-increase-headroom.lua")

    print_header("Audio installed successfully! Reboot to finish setup.")