#!/usr/bin/env python3
"""
Persistent broker: holds one live VM console session across tool calls.
- Retries the rate-limited port until it gets a live VM.
- Waits for the BusyBox shell.
- Serves a Unix socket (broker.sock); bytes sent there go to the VM shell.
- All VM output is appended raw to session.log.
- Sends keepalive newlines when idle.
"""
import socket, time, select, os, sys, signal

HOST = os.environ.get("BROKER_HOST", "x.x.x.x")
PORT = int(os.environ.get("BROKER_PORT", "30767"))
BASE = os.path.dirname(os.path.abspath(__file__))
LOGF = os.path.join(BASE, "session.log")
STATEF = os.path.join(BASE, "broker.state")
SOCKF = os.path.join(BASE, "broker.sock")

logf = open(LOGF, "ab", buffering=0)

def log(b):
    try:
        logf.write(b)
    except Exception:
        pass

def state(s):
    try:
        with open(STATEF, "w") as f:
            f.write(s + "\n")
    except Exception:
        pass

def note(s):
    ts = time.strftime("%H:%M:%S", time.gmtime())
    log(b"\n### [" + ts.encode() + b"] " + s.encode(errors="replace") + b"\n")
    state(s)

def connect_live():
    attempt = 0
    while True:
        attempt += 1
        try:
            s = socket.create_connection((HOST, PORT), timeout=5)
        except Exception as e:
            if attempt % 5 == 0:
                note(f"connect err {e!r} (attempt {attempt})")
            time.sleep(5)
            continue
        s.setblocking(False)
        time.sleep(0.5)
        r, _, _ = select.select([s], [], [], 0.3)
        if r:
            try:
                d = s.recv(4096)
            except Exception:
                d = b""
            if d == b"":
                if attempt % 5 == 0:
                    note(f"rejected (attempt {attempt})")
                s.close()
                time.sleep(5)
                continue
            note(f"LIVE after {attempt} attempts, banner={d!r}")
            return s
        note(f"LIVE (silent) after {attempt} attempts")
        return s

def wait_shell(s):
    end = time.time() + 120
    last = time.time()
    buf = b""
    while time.time() < end:
        r, _, _ = select.select([s], [], [], 0.5)
        if r:
            try:
                d = s.recv(65536)
            except Exception:
                return False
            if d == b"":
                return False
            log(d)
            buf += d
            if b"BusyBox v1.37.0" in buf or b"~ $" in buf:
                return True
        if time.time() - last > 4:
            try:
                s.sendall(b"\n")
            except Exception:
                return False
            last = time.time()
    return False

def serve(s):
    try:
        os.unlink(SOCKF)
    except FileNotFoundError:
        pass
    srv = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    srv.bind(SOCKF)
    os.chmod(SOCKF, 0o777)
    srv.listen(16)
    srv.setblocking(False)
    clients = []
    last_traffic = time.time()
    note("SHELL READY / relay up")
    while True:
        try:
            r, _, _ = select.select([s, srv] + clients, [], [], 0.5)
        except Exception:
            break
        if s in r:
            try:
                d = s.recv(65536)
            except Exception as e:
                note(f"vm recv err {e!r}")
                break
            if d == b"":
                note("VM EOF")
                break
            log(d)
            for c in clients[:]:
                try:
                    c.sendall(d)
                except Exception:
                    clients.remove(c)
            last_traffic = time.time()
        if srv in r:
            try:
                c, _ = srv.accept()
                c.setblocking(False)
                clients.append(c)
            except Exception:
                pass
        for c in clients[:]:
            if c in r:
                try:
                    d = c.recv(65536)
                except Exception:
                    clients.remove(c)
                    continue
                if d == b"":
                    clients.remove(c)
                    try: c.close()
                    except Exception: pass
                    continue
                log(b"\n<<< " + d + b"\n")
                try:
                    s.sendall(d)
                except Exception as e:
                    note(f"vm send err {e!r}")
                    break
                last_traffic = time.time()
        if time.time() - last_traffic > 15:
            try:
                s.sendall(b"\n")
            except Exception:
                note("heartbeat send err")
                break
            last_traffic = time.time()
    try: srv.close()
    except Exception: pass
    for c in clients:
        try: c.close()
        except Exception: pass
    try: os.unlink(SOCKF)
    except Exception: pass

def main():
    signal.signal(signal.SIGTERM, lambda *a: sys.exit(0))
    note("broker start")
    while True:
        s = connect_live()
        if not wait_shell(s):
            note("shell wait failed; reconnecting")
            try: s.close()
            except Exception: pass
            time.sleep(3)
            continue
        serve(s)
        try: s.close()
        except Exception: pass
        note("session ended; reconnecting")
        time.sleep(3)

if __name__ == "__main__":
    main()