"""The load-bearing line, isolated: one verifier, one line toggled, nothing
else changed. The allowlist, the constant-time compare, the RSA check and
the expiry check are identical on both rows.

    python3 oneline.py
"""

import hmac

import jwt

NOW = 1_700_000_000
HONEST = {"sub": "alice", "role": "user", "exp": NOW + 3600}
GREEDY = {"sub": "alice", "role": "admin", "exp": NOW + 86400 * 365}

honest_token = jwt.encode(HONEST, "RS256", jwt.PRIVATE_KEY)
forged = jwt.encode(GREEDY, "HS256", jwt.PUBLIC_KEY)   # signed with the PUBLIC key


def verify(token, key, server_alg, trust_the_token):
    header, payload, sig, signing_input = jwt.split(token)
    if trust_the_token:
        alg = header["alg"]           # LINE A: the token decides
    else:
        alg = server_alg              # LINE B: the server decides
    if alg not in ("HS256", "RS256"):
        return None
    if alg == "HS256":
        ok = hmac.compare_digest(jwt.hs256(signing_input, key), sig)
    else:
        ok = jwt.rs256_check(signing_input, sig, key)
    if not ok or payload.get("exp", 0) < NOW:
        return None
    return payload


def role(result):
    return result["role"] if result else "REJECTED"


for trust, label in ((True, 'alg = header["alg"]  '), (False, "alg = server_alg     ")):
    h = verify(honest_token, jwt.PUBLIC_KEY, "RS256", trust)
    f = verify(forged, jwt.PUBLIC_KEY, "RS256", trust)
    print(f"  {label}  honest -> {role(h):<9} forged -> {role(f)}")