#!/usr/bin/env python3
# Per-COMMAND auth variant (model A: "authenticate each time").
# Same daemon/ops as the persistent-connection harness, but every command
# opens a FRESH connection, does a full signed-HELLO R4 handshake, sends the
# command, reads the reply, then closes. This is what cubec one-shot /
# stress.sh does. Unlike stress.sh, THIS harness tracks per-command errors
# so we can compare the error rate head-to-head against auth-once.
import os, sys, socket, struct, time, threading, csv, hmac, hashlib, subprocess, tempfile, shutil
from collections import defaultdict
import statistics as st

REPO = os.environ.get("REPO", "/home/CUBELinux/CUBELinux-2")
SRV = os.path.join(REPO, "target/release/cube-server")
if not os.path.exists(SRV):
    sys.exit("cube-server release binary missing; build first")

USERS = int(os.environ.get("USERS", "8"))
DURATION = int(os.environ.get("DURATION", "150"))
PSK = b"lan-shared-key-1234"

WORK = "/root/cube-stress"
os.makedirs(WORK, exist_ok=True)
TMPD = tempfile.mkdtemp(prefix=f"{WORK}/run-percmd-")
STORE = os.path.join(TMPD, f"store-{os.getpid()}.json")
PSK_FILE = os.path.join(TMPD, "psk.txt")
open(PSK_FILE, "w").write(PSK.decode())
SOCK = os.path.join(TMPD, "cube.sock")
RUN_LOG = os.path.join(TMPD, "run.log")
DAEMON_LOG = os.path.join(TMPD, "daemon.log")
ERR_LOG = os.path.join(TMPD, "run-errors.log")
CSV = os.path.join(TMPD, "cmds.csv")

TENANTS = ["alpha","bravo","charlie","delta","echo","foxtrot","golf","hotel","india","juliet"]
OWNERS  = ["alice","bob","carol","dave","erin","frank","grace","heidi","ivan","judy"]

def log(*a):
    with open(RUN_LOG, "a") as f:
        f.write(" ".join(map(str,a)) + "\n")
log(f"USERS={USERS} DURATION={DURATION}s MODE=per-command-auth (fresh handshake every op)")

def write_frame(sock, payload):
    b = payload.encode(); sock.sendall(struct.pack("<I", len(b)) + b)
def read_frame(sock):
    n = struct.unpack("<I", sock.recv(4))[0]
    return sock.recv(n).decode()
def sign_hello(psk, nonce, tenant, owner_local, owner_remote=""):
    msg = f"{nonce}|{tenant}|{owner_local}|{owner_remote}".encode()
    return hmac.new(psk, msg, hashlib.sha256).hexdigest()

daemon_args = [SRV, "--socket", SOCK, "--store", STORE, "--auth-key", PSK_FILE]
dp = subprocess.Popen(daemon_args, stdout=open(DAEMON_LOG,"w"), stderr=subprocess.STDOUT)
for _ in range(15):
    if os.path.exists(SOCK): break
    time.sleep(1)
if not os.path.exists(SOCK):
    log("daemon failed to start"); sys.exit(1)
log(f"daemon up pid={dp.pid} on {SOCK} [auth REQUIRED, per-command handshake]")

# metrics
hs_results = []          # (ok, ms) per handshake
results = []             # (user, op, ok, latency_ns)
hs_fail = 0
stop = threading.Event()
csv_lock = threading.Lock()

def make_cmd(u, c):
    pair = c // 8
    X = ((pair + u*7) % 200) + 1
    Y = ((pair*3 + u*5) % 200) + 1
    Z = ((pair*11 + u*13) % 200) + 1
    W = ((pair*17 + u*3) % 200) + 1
    p = f"/c{W:03d}/z{X:03d}/y{Y:03d}/x{Z:03d}"
    op = c % 8
    if op == 0:  return op, f"prog {p} const {X} halt"
    if op == 1:  return op, f"run {p}"
    if op == 2:  return op, f"grant {OWNERS[u]} r 100.1.1.1"
    if op == 3:  return op, f"revoke {OWNERS[u]} 100.1.1.1"
    if op == 4:  return op, "query fn"
    if op == 5:  return op, "stats"
    if op == 6:
        if os.environ.get("NO_AUDIT") == "1":
            return op, "stats"   # drop the slow audit op to replicate legacy op mix
        return op, "audit"
    return op, "stats"

def one_cmd(u):
    """Open a fresh connection, handshake, run ONE command, close. Returns (hs_ok, hs_ms, op_ok, op_ms)."""
    tenant, owner = TENANTS[u], OWNERS[u]
    s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    try:
        s.connect(SOCK)
    except OSError:
        return (0, 0.0, 0, 0.0)
    s.settimeout(3.0)
    t0 = time.time()
    try:
        chal = read_frame(s)
        if not chal.lower().startswith("challenge "):
            s.close(); return (0, 0.0, 0, 0.0)
        nonce = chal[len("challenge "):].strip()
        sig = sign_hello(PSK, nonce, tenant, owner)
        write_frame(s, f"HELLO {tenant} {owner} {sig}")
        ack = read_frame(s)
        hs_ok = 1 if ack.startswith("ok:") else 0
        hs_ms = (time.time() - t0) * 1000.0
    except Exception:
        s.close(); return (0, 0.0, 0, 0.0)
    if hs_ok == 0:
        s.close(); return (0, hs_ms, 0, 0.0)
    # now run one command on this fresh connection
    op, cmd = make_cmd(u, one_cmd.counter[0])
    one_cmd.counter[0] += 1
    t1 = time.time()
    try:
        write_frame(s, cmd); reply = read_frame(s)
        ok = 0 if reply.lower().startswith("error") else 1
    except socket.timeout:
        ok = 0; reply = "TIMEOUT(>3s)"
    except Exception as e:
        ok = 0; reply = f"EXC:{e}"
    op_ms = (time.time() - t1) * 1e9
    s.close()
    return (1, hs_ms, ok, op_ms)

def drive_user(u):
    while not stop.is_set() and time.time() < END:
        ho, hms, ok, oms = one_cmd(u)
        with csv_lock:
            hs_results.append((ho, hms))
            results.append((u, -2, ok, oms))   # op -2 = a per-cmd auth+op unit

one_cmd.counter = [0]
END = time.time() + DURATION
threads = [threading.Thread(target=drive_user, args=(u,), daemon=True) for u in range(USERS)]
for t in threads: t.start()
for t in threads: t.join()
stop.set()
dp.terminate()
try: dp.wait(timeout=5)
except Exception: dp.kill()

with open(CSV, "w", newline="") as f:
    w = csv.writer(f); w.writerow(["user","op","ok","latency_ns"])
    for r in results: w.writerow(r)

hs_ok = sum(1 for h,_ in hs_results if h == 1)
ops = [r for r in results if r[1] == -2]
ok = sum(1 for r in ops if r[2] == 1)
lat = [r[3] for r in ops]
def pct(xs,p):
    xs=sorted(xs); k=int(round((p/100)*(len(xs)-1))); return xs[k]/1e6
log("===== AGGREGATE RESULTS (per-command auth, model A) =====")
log(f"users                : {USERS}")
log(f"handshakes           : {len(hs_results)}")
log(f"handshake ok         : {hs_ok}  ({100*hs_ok/max(len(hs_results),1):.2f}%)")
if hs_results:
    hms=[m for _,m in hs_results]
    log(f"handshake ms mean/max: {st.mean(hms):.3f} / {max(hms):.3f}")
log(f"total ops            : {len(ops)}  (incl. their inline handshake)")
log(f"successful (ok=1)    : {ok}  ({100*ok/max(len(ops),1):.2f}%)")
log(f"failed (ok=0)        : {len(ops)-ok}")
if lat:
    log(f"op latency ms mean/med : {st.mean(lat)/1e6:.3f} / {st.median(lat)/1e6:.3f}")
    log(f"op latency ms p95/p99  : {pct(lat,95):.3f} / {pct(lat,99):.3f}")
    log(f"op latency ms max      : {max(lat)/1e6:.3f}")
log(f"TMPD={TMPD}")
log("DONE.")
print(f"DONE. results in {TMPD}")