#!/usr/bin/python3
"""SpanOPS Station setup helper: the only part of Station that runs as root.

Started by Station through pkexec (polkit action app.spanops.station.setup,
administrator authentication). It does exactly two things, with strictly checked input:

    station-setup-helper add NAME URI [use]
                                         add a label printer to CUPS (URI: socket://HOST[:PORT], or a
                                         usb:// device CUPS detects); "use" also makes Station use it
    station-setup-helper use NAME        make NAME the printer Station uses on this computer
                                         ([printer] queue in /etc/spanops/station.conf)

Prints one line ("ok: ..." or "error: ...") and exits 0 on success.
"""
import ipaddress
import os
import re
import subprocess
import sys
import tempfile

CONF = "/etc/spanops/station.conf"
ZPL_DRIVER = "drv:///sample.drv/zebra.ppd"  # CUPS' "Zebra ZPL Label Printer"; Station sends raw ZPL
NAME = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]{0,63}$")
HOST = re.compile(r"^(?=.{1,253}$)[A-Za-z0-9]([A-Za-z0-9-]{0,62}[A-Za-z0-9])?(\.[A-Za-z0-9]([A-Za-z0-9-]{0,62}[A-Za-z0-9])?)*$")
ENV = {"PATH": "/usr/sbin:/usr/bin:/sbin:/bin", "LC_ALL": "C"}


def fail(msg: str) -> None:
    print(f"error: {msg}")
    sys.exit(1)


def run(args: list[str], timeout: int = 30) -> str:
    p = subprocess.run(args, capture_output=True, text=True, timeout=timeout, env=ENV)
    if p.returncode != 0:
        fail((p.stderr or p.stdout).strip().splitlines()[-1] if (p.stderr or p.stdout).strip() else
             f"{args[0]} failed ({p.returncode})")
    return p.stdout


def check_uri(uri: str) -> str:
    if uri.startswith("socket://"):
        rest = uri[len("socket://"):]
        if rest.count(":") > 1 or "/" in rest or "@" in rest:
            fail("use an IPv4 address or a host name, optionally with :port")
        host, _, port = rest.rpartition(":") if rest.count(":") == 1 else (rest, "", "9100")
        try:
            ipaddress.ip_address(host)
        except ValueError:
            if not HOST.match(host):
                fail("that isn't a valid network address")
        if not port.isdigit() or not 1 <= int(port) <= 65535:
            fail("that isn't a valid port")
        return f"socket://{host}:{int(port)}"
    if uri.startswith("usb://"):
        detected = run(["lpinfo", "--include-schemes", "usb", "-v"], timeout=60).split()
        if uri not in detected:
            fail("that USB printer isn't connected to this computer any more")
        return uri
    fail("only network (socket://) and USB (usb://) printers can be added")


def add(name: str, uri: str) -> None:
    if not NAME.match(name):
        fail("a printer name may use letters, digits, dot, dash and underscore (up to 64)")
    if name in run(["lpstat", "-e"]).split():
        fail(f"there is already a printer called {name}")
    uri = check_uri(uri)
    run(["lpadmin", "-p", name, "-E", "-v", uri, "-m", ZPL_DRIVER,
         "-D", "SpanOPS label printer"])
    run(["cupsenable", name])
    run(["cupsaccept", name])
    print(f"ok: added {name} ({uri})")


def use(name: str) -> None:
    if not NAME.match(name) or name not in run(["lpstat", "-e"]).split():
        fail(f"there is no printer called {name} on this computer")
    try:
        with open(CONF, encoding="utf-8") as f:
            lines = f.read().splitlines()
    except FileNotFoundError:
        fail(f"{CONF} is missing: reinstall SpanOPS Station")
    out, section, done = [], "", False
    for line in lines:
        s = line.strip()
        if s.startswith("[") and s.endswith("]"):
            if section == "printer" and not done:
                out.append(f"queue = {name}")
                done = True
            section = s[1:-1].strip().lower()
        elif section == "printer" and re.match(r"^\s*queue\s*[=:]", line):
            if not done:
                out.append(f"queue = {name}")
                done = True
            continue  # drop the old line (and any duplicate)
        out.append(line)
    if not done:
        if section != "printer":
            out += ["", "[printer]"]
        out.append(f"queue = {name}")
    st = os.stat(CONF)
    fd, tmp = tempfile.mkstemp(dir=os.path.dirname(CONF), prefix=".station.conf.")
    try:
        with os.fdopen(fd, "w", encoding="utf-8") as f:
            f.write("\n".join(out) + "\n")
        os.chown(tmp, st.st_uid, st.st_gid)
        os.chmod(tmp, st.st_mode & 0o7777)
        os.replace(tmp, CONF)
    except BaseException:
        if os.path.exists(tmp):
            os.unlink(tmp)
        raise
    print(f"ok: Station uses {name}")


def main() -> None:
    if os.geteuid() != 0:
        fail("must run as root (through pkexec)")
    a = sys.argv[1:]
    if len(a) in (3, 4) and a[0] == "add" and a[3:] in ([], ["use"]):
        add(a[1], a[2])
        if a[3:] == ["use"]:
            use(a[1])
    elif len(a) == 2 and a[0] == "use":
        use(a[1])
    else:
        fail("usage: add NAME URI [use] | use NAME")


if __name__ == "__main__":
    main()
