"""A JSON Web Token, signed and verified by hand. No library.

A JWT is three base64url segments joined by dots:

    base64url(header) . base64url(payload) . base64url(signature)

and the signature covers the first two segments *as text*, including the
dots. The header is a JSON object. One of its fields is `alg`, and `alg`
names the algorithm the verifier should use.

That last sentence is the toy. The header is a document the attacker hands
you, and a verifier that reads `alg` out of it is asking the attacker how the
attacker should be checked. Four verifiers here, in the order a codebase
usually acquires them:

    verify_naive      dispatch on header["alg"]
    verify_blocklist  ... but refuse "none"
    verify_allowlist  ... but only allow HS256 and RS256
    verify_strict     never read header["alg"] at all

The first three all accept a forged token. See commentary section 6.

WARNING, and it is not a formality: the RS256 here is *textbook* RSA --
raw modular exponentiation of a bare SHA-256 digest, no PKCS#1 padding, a
511-bit modulus, two primes you can read on screen. It is insecure by
construction and exists only so this toy has an algorithm whose verification
key is genuinely public rather than asserted to be. Do not copy it.
"""

import base64
import hashlib
import hmac
import json

# ------------------------------------------------------------- base64url

def b64e(raw):
    """base64url with the '=' padding stripped, per RFC 7515 appendix C."""
    return base64.urlsafe_b64encode(raw).rstrip(b"=").decode()


def b64d(txt):
    """The inverse. Padding is recomputed, never transmitted."""
    return base64.urlsafe_b64decode(txt + "=" * (-len(txt) % 4))


def _json(obj):
    """Canonical JSON: sorted keys, no whitespace. Not required by the spec
    -- it is required by this page, so every token below is byte-identical
    on every run.
    """
    return json.dumps(obj, sort_keys=True, separators=(",", ":")).encode()


# ---------------------------------------------------------- the key pair

# A 511-bit modulus from two 256-bit primes, found once by a Miller-Rabin
# search upward from 2**255 and then hardcoded, so there is no key
# generation and no randomness anywhere in this toy.
P = 57896044618658097711785492504343953926634992332820282019728792003956564820063
Q = 57896044618658097711785492504343953926634992332820282019828792003956564820137
N = P * Q
E = 65537
D = pow(E, -1, (P - 1) * (Q - 1))

# PUBLIC_KEY is what the server holds in order to *verify* RS256, and what
# anybody at all can fetch from a /.well-known/jwks.json. It is bytes,
# because that is how a real verifier receives it: a PEM blob, a JWKS entry,
# a file read off disk. Section 6.3 is about what happens when those bytes
# reach a function that was expecting an HMAC secret.
PUBLIC_KEY = _json({"kty": "RSA", "n": N, "e": E})
PRIVATE_KEY = _json({"kty": "RSA", "n": N, "d": D})

# The shared secret for a symmetric deployment. Section 6.7 uses it to show
# where the attack stops working.
SECRET = b"correct horse battery staple"


# ------------------------------------------------------------ primitives

def hs256(signing_input, key):
    """HMAC-SHA-256. `key` is arbitrary bytes -- any bytes at all, which is
    exactly the property section 6.3 turns into a forgery.
    """
    return hmac.new(key, signing_input, hashlib.sha256).digest()


def rs256_sign(signing_input, priv):
    """Textbook RSA: s = h^d mod n, over the bare digest. Deterministic,
    unpadded, and insecure -- see the module docstring.
    """
    k = json.loads(priv)
    h = int.from_bytes(hashlib.sha256(signing_input).digest(), "big")
    return pow(h, k["d"], k["n"]).to_bytes(64, "big")


def rs256_check(signing_input, sig, pub):
    """s^e mod n == h. Note the asymmetry that makes RS256 worth having:
    this needs only `pub`, while rs256_sign needs `priv`.
    """
    k = json.loads(pub)
    h = int.from_bytes(hashlib.sha256(signing_input).digest(), "big")
    return pow(int.from_bytes(sig, "big"), k["e"], k["n"]) == h


# ----------------------------------------------------------- encode side

def encode(payload, alg, key, extra_header=None):
    """Mint a token. Anyone can call this, including an attacker -- the
    only question is which `key` they have, and section 6.3 is about a `key`
    they are given on purpose.
    """
    header = {"alg": alg, "typ": "JWT"}
    if extra_header:
        header.update(extra_header)
    h, p = b64e(_json(header)), b64e(_json(payload))
    signing_input = f"{h}.{p}".encode()
    if alg == "HS256":
        sig = hs256(signing_input, key)
    elif alg == "RS256":
        sig = rs256_sign(signing_input, key)
    elif alg == "none":
        sig = b""
    else:
        raise ValueError(alg)
    return f"{h}.{p}.{b64e(sig)}"


def split(token):
    """(header, payload, signature bytes, signing input). The signing input
    is the first two segments and the dot between them, as ASCII -- not the
    decoded JSON, which is why re-serialising a header breaks a signature.
    """
    h, p, s = token.split(".")
    return json.loads(b64d(h)), json.loads(b64d(p)), b64d(s), f"{h}.{p}".encode()


def peek(token):
    """Header and payload, with no key and no signature check. This is not
    a weakness in the toy; it is what base64 is. Section 6.6.
    """
    h, p, _s = token.split(".")
    return json.loads(b64d(h)), json.loads(b64d(p))


# --------------------------------------------------------- the verifiers

def verify_naive(token, key, now=0):
    """v0. Asks the token how the token should be checked.

    Returns the payload on success, None on rejection. Read the first two
    lines of the body: `alg` comes out of the header, and the header came
    from whoever sent the token.
    """
    header, payload, sig, signing_input = split(token)
    alg = header["alg"]
    if alg == "none":
        ok = sig == b""
    elif alg == "HS256":
        ok = hmac.compare_digest(hs256(signing_input, key), sig)
    elif alg == "RS256":
        ok = rs256_check(signing_input, sig, key)
    else:
        return None
    if not ok:
        return None
    if "exp" in payload and payload["exp"] < now:
        return None
    return payload


def verify_blocklist(token, key, now=0):
    """v1. The reflex fix, and the one the industry shipped in 2015:
    refuse the algorithm that turned out to be a hole.
    """
    header, _payload, _sig, _si = split(token)
    if header["alg"] == "none":
        return None
    return verify_naive(token, key, now)


def verify_allowlist(token, key, now=0, allowed=("HS256", "RS256")):
    """v2. The careful fix. Name the algorithms you are prepared to accept,
    up front, and reject everything else. Both of these are real, strong,
    standard algorithms. This verifier is still broken -- section 6.3.
    """
    header, _payload, _sig, _si = split(token)
    if header["alg"] not in allowed:
        return None
    return verify_naive(token, key, now)


def verify_strict(token, key, alg, now=0):
    """v3. The caller names the algorithm. `header["alg"]` is never read --
    the header is now signed *data*, not an instruction.

    `alg` must be the algorithm that matches `key`. Passing "HS256" here
    together with an RSA public key reproduces the whole bug (section 7.2),
    so the pairing is the fix, not the keyword.
    """
    _header, payload, sig, signing_input = split(token)
    if alg == "HS256":
        ok = hmac.compare_digest(hs256(signing_input, key), sig)
    elif alg == "RS256":
        ok = rs256_check(signing_input, sig, key)
    else:
        raise ValueError(alg)
    if not ok:
        return None
    if payload.get("exp", 0) < now:
        return None
    return payload