dotfiles/os/wireless/import.py

#!/usr/bin/env nix-shell
#! nix-shell -i python3
#! nix-shell -p python3 python3Packages.pyaml

"""
Exports Wi-Fi networks configuration stored in pass into a format readable by Nix.
"""

# TODO EAP ca_cert=/etc/ssl/... probably won't work. Example fix:
# builtins.fetchurl {
#     url = "https://letsencrypt.org/certs/isrgrootx1.pem";
#     sha256 = "sha256:1la36n2f31j9s03v847ig6ny9lr875q3g7smnq33dcsmf2i5gd92";
# }

import hashlib
import json
import os
import subprocess

import yaml

# passpy doesn't handle encoding properly, so doing this with calls

PASSWORD_STORE = os.path.expanduser("~/.local/share/pass")
SUBFOLDER = "wifi"
SEPARATE_PASSWORDS = False
# TODO Find a way to make then env file available at whatever time it is needed


class Password:
    all: list["Password"] = list()

    def __init__(self, path: str, content: str):
        self.path = path
        self.content = content

        Password.all.append(self)

    def var(self) -> str:
        # return self.path.split("/")[-1].upper()
        m = hashlib.sha256()
        m.update(self.path.encode())
        return m.hexdigest().upper()

    def val(self) -> str:
        return self.content

    def exists(self) -> bool:
        return not not self.content

    def key(self) -> str:
        if SEPARATE_PASSWORDS:
            return f"@{self.var()}@"
        else:
            return self.val()

    @classmethod
    def vars(cls) -> dict[str, str]:
        vars = dict()
        for password in cls.all:
            if not password.content:
                continue
            var = password.var()
            assert var not in vars, f"Duplicate key: {var}"
            vars[var] = password.val()
        return vars


def list_networks() -> list[str]:
    paths = []
    pass_folder = os.path.join(PASSWORD_STORE, SUBFOLDER)
    for filename in os.listdir(pass_folder):
        if not filename.endswith(".gpg"):
            continue
        filepath = os.path.join(pass_folder, filename)
        if not os.path.isfile(filepath):
            continue

        file = filename[:-4]
        path = os.path.join(SUBFOLDER, file)
        paths.append(path)
    return paths


def format_wpa_supplicant_conf(conf: dict, indent: str = "") -> str:
    lines = []
    for k, v in conf.items():
        if isinstance(v, str):
            val = '"' + v.replace('"', '\\"') + '"'
        elif isinstance(v, Password):
            val = v.key()
        elif isinstance(v, list):
            assert all(
                map(lambda i: isinstance(i, str), v)
            ), "Only list of strings supported"
            val = " ".join(v)
        else:
            val = str(v)
        lines.append(f"{indent}{k}={val}")
    return "\n".join(lines)


networks = {}
for path in list_networks():
    proc = subprocess.run(["pass", path], stdout=subprocess.PIPE)
    proc.check_returncode()

    raw = proc.stdout.decode()
    split = raw.split("\n")

    password = Password(path, split[0])
    data = yaml.safe_load("\n".join(split[1:])) or dict()
    # print(path, data)  # DEBUG

    # Helpers to prevent repetition
    suffixes = data.pop("suffixes", [""])
    data.setdefault("key_mgmt", ["WPA-PSK"] if password.exists() else ["NONE"])
    if password:
        if any(map(lambda m: "PSK" in m.split("-"), data["key_mgmt"])):
            data["psk"] = password
        if "NONE" in data["key_mgmt"]:
            data["wep_key0"] = password
        if any(map(lambda m: "EAP" in m.split("-"), data["key_mgmt"])):
            data["password"] = password
    assert "ssid" in data, f"{path}: Missing SSID"

    # # Output wpa_supplicant conf, for debug
    # for suffix in suffixes:
    #     wpas = data.copy()
    #     wpas["ssid"] += suffix
    #     print(f"# {path}")
    #     print("network={")
    #     print(format_wpa_supplicant_conf(wpas, indent="    "))
    #     print("}")
    #     print()

    # Convert to nix configuration
    ssid = data.pop("ssid")
    network = {}
    key_mgmt = data.pop("key_mgmt", None)
    psk = data.pop("psk", None)
    priority = data.pop("priority", None)
    # No support for hidden
    # No support for extraConfig (all is assumed to be auth)
    if key_mgmt:
        network["authProtocols"] = key_mgmt
    if psk:
        network["psk"] = psk.key()
        if data:
            raise NotImplementedError(f"{path}: Unhandled non-auth extra: {data}")
    else:
        if data:
            network["auth"] = format_wpa_supplicant_conf(data)
    if priority:
        network["priority"] = int(priority)

    for suffix in suffixes:
        networks[ssid + suffix] = network

with open("networks.json", "w") as fd:
    json.dump(networks, fd, indent=4)

with open("networks.env", "w") as fd:
    if SEPARATE_PASSWORDS:
        for k, v in Password.vars().items():
            print(f"{k}={v}", file=fd)