#!/usr/bin/env python3
"""Capture-first LSX diagnostic for FIFA 17 on 127.0.0.1:4216.

Goal: learn the REAL handshake the game speaks, losing nothing to a crash.
- Detects who sends first (server-initiates vs client-initiates).
- Sends our best-guess Challenge, then captures the client's ChallengeResponse
  (its key + exact XML format) -- FLUSHED TO DISK BEFORE we send anything risky.
- Then runs the full instrumented handshake (H, session key, encrypted loop),
  logging every computed value + every frame both directions.
So even if the game crashes on our ChallengeAccepted, the key capture is saved.
"""
import socket, time, threading, os, sys
import lsx_responder as L   # reuse crypto + build_reply; importing does NOT run main

LOG = "/tmp/lsx_capture.log"
RAWDIR = "/tmp/lsx_raw"
os.makedirs(RAWDIR, exist_ok=True)
_n = 0

def log(m):
    line = "[%s] %s" % (time.strftime("%H:%M:%S"), m)
    print(line, flush=True)
    with open(LOG, "a") as f:
        f.write(line + "\n"); f.flush(); os.fsync(f.fileno())

def dump(tag, b):
    log("%s : %dB" % (tag, len(b)))
    for i in range(0, min(len(b), 512), 16):
        c = b[i:i+16]
        hx = " ".join("%02x" % x for x in c)
        asc = "".join(chr(x) if 32 <= x < 127 else "." for x in c)
        log("    %04x: %-47s %s" % (i, hx, asc))
    try:
        p = "%s/%s_%d.bin" % (RAWDIR, tag.replace(" ", "_").replace("#", "").replace(":", ""), int(time.time()*1000) % 1000000)
        open(p, "wb").write(b)
    except Exception:
        pass

def serve(conn, addr):
    global _n; _n += 1; n = _n
    log("================= CONN #%d from %s =================" % (n, addr))
    # Phase 0 -- does the CLIENT speak first? (would mean lsx_responder has the
    # initiator backwards, a prime crash suspect)
    conn.settimeout(3.0)
    try:
        pre = conn.recv(8192)
        if pre:
            dump("C%d PRE (client spoke FIRST)" % n, pre)
            log("C%d: !!! client initiates -- our server-sends-Challenge model is WRONG" % n)
    except socket.timeout:
        log("C%d: silent 3s -> server initiates (send Challenge)" % n)
        pre = b""
    except Exception as e:
        log("C%d pre-recv err: %s" % (n, e)); pre = b""

    # Phase 1 -- send our best-guess plaintext Challenge, capture the reply.
    chal = ('<LSX><Event sender="EALS"><Challenge key="%s" build="%s" version="%s"/>'
            '</Event></LSX>' % (L.CHALLENGE_KEY, L.BUILD, L.VERSION))
    try:
        conn.sendall(chal.encode() + b"\0")
        log("C%d >> Challenge sent (%dB): %s" % (n, len(chal)+1, chal))
    except Exception as e:
        log("C%d send-Challenge err: %s" % (n, e)); return

    try:
        rep = conn.recv(8192)
    except Exception as e:
        log("C%d recv-after-Challenge err: %s" % (n, e)); return
    if not rep:
        log("C%d: client closed after our Challenge (no ChallengeResponse) -- "
            "Challenge format likely rejected" % n); return
    dump("C%d CHALLENGE-RESPONSE (client)" % n, rep)   # <-- THE KEY CAPTURE, already flushed

    # Phase 2 -- log what WE would compute (do NOT let a crypto error abort logging)
    txt = rep.split(b"\0")[0].decode(errors="replace")
    log("C%d client-reply text: %s" % (n, txt))
    import re
    mk = re.search(r'key="([^"]*)"', txt)
    mr = re.search(r'response="([^"]*)"', txt)
    client_key = mk.group(1) if mk else L.CHALLENGE_KEY
    client_resp = mr.group(1) if mr else None
    log("C%d parsed: client_key=%r  client_response=%r" % (n, client_key, client_resp))
    try:
        our_h = L.challenge_response(client_key)
        log("C%d our computed H (ChallengeAccepted.response) = %s" % (n, our_h))
        if client_resp:
            log("C%d MATCH client_response==our_H ? %s" % (n, client_resp == our_h))
        skey = L.derive_session_key(our_h)
        log("C%d derived session_key = %s" % (n, skey.hex()))
    except Exception as e:
        log("C%d crypto compute err: %s" % (n, e)); our_h = None; skey = None

    # Phase 3 -- OPTIONAL: send ChallengeAccepted + run the encrypted loop.
    # Guarded by env so the first run can stay capture-only (no risky send).
    if os.environ.get("LSX_FULL") == "1" and our_h and skey:
        try:
            conn.sendall(L.resp(1, 'ChallengeAccepted response="%s"' % our_h, "EALS").encode() + b"\0")
            log("C%d >> ChallengeAccepted sent" % n)
        except Exception as e:
            log("C%d send-Accepted err: %s" % (n, e)); return
        while True:
            try:
                data = conn.recv(65536)
            except Exception as e:
                log("C%d recv-loop err: %s" % (n, e)); break
            if not data:
                log("C%d: client closed" % n); break
            dump("C%d ENC-IN" % n, data)
            for chunk in filter(None, data.split(b"\0")):
                try:
                    xml = L.lsx_decrypt(chunk + b"\0", skey)
                    log("C%d << decrypted: %s" % (n, xml))
                except Exception as e:
                    log("C%d decrypt fail: %s (raw %s)" % (n, e, chunk[:40])); continue
                mm = L.REQ_RE.search(xml)
                if mm:
                    reply = L.build_reply(mm.group(1), mm.group(2),
                                          dict(L.ATTR_RE.findall(mm.group(3))))
                    try:
                        conn.sendall(L.lsx_encrypt(reply, skey))
                        log("C%d >> %s" % (n, reply))
                    except Exception as e:
                        log("C%d send-reply err: %s" % (n, e)); break
    else:
        log("C%d: capture-only (set LSX_FULL=1 to attempt full handshake). Holding 10s." % n)
        time.sleep(10)
    try: conn.close()
    except Exception: pass

def main():
    s = socket.socket()
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    s.bind(("127.0.0.1", 4216)); s.listen(8)
    log("=== capture_lsx listening on 127.0.0.1:4216  (LSX_FULL=%s) ===" %
        os.environ.get("LSX_FULL", "0"))
    while True:
        c, a = s.accept()
        threading.Thread(target=serve, args=(c, a), daemon=True).start()

if __name__ == "__main__":
    main()