"""Stdlib-only tests (no pytest): plain asserts in functions called from a
__main__ block. Run: `python3 test_jwt.py`.

The headline tests pin the two forgeries the commentary is about -- the
`alg: none` token (section 6.2) and the HS256/RS256 key-confusion token
(section 6.3) -- pin which verifier stops which, and pin the two boundaries
(sections 6.7 and 7.2). The rest are round-trips, so a "fix" that simply
rejects everything cannot pass.
"""

import json

import jwt

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

RS_TOKEN = jwt.encode(HONEST, "RS256", jwt.PRIVATE_KEY)
HS_TOKEN = jwt.encode(HONEST, "HS256", jwt.SECRET)
F_NONE = jwt.encode(GREEDY, "none", None)
F_CONF = jwt.encode(GREEDY, "HS256", jwt.PUBLIC_KEY)


def _all(token, key=None):
    """(v0, v1, v2, v3) against an RS256 deployment unless told otherwise."""
    key = jwt.PUBLIC_KEY if key is None else key
    return (jwt.verify_naive(token, key, NOW),
            jwt.verify_blocklist(token, key, NOW),
            jwt.verify_allowlist(token, key, NOW),
            jwt.verify_strict(token, key, "RS256", NOW))


def test_honest_tokens_round_trip():
    """Every verifier accepts the genuine token, and reads the real claims."""
    for result in _all(RS_TOKEN):
        assert result == HONEST, result
    assert jwt.verify_strict(HS_TOKEN, jwt.SECRET, "HS256", NOW) == HONEST


def test_tampering_the_payload_is_caught():
    """The recorded aha, and it holds: one byte, and the signature dies."""
    h, p, s = RS_TOKEN.split(".")
    one_byte = jwt.b64e(jwt.b64d(p).replace(b"alice", b"alicf"))
    escalated = jwt.b64e(jwt.b64d(p).replace(b'"user"', b'"admin"'))
    for mutant in (f"{h}.{one_byte}.{s}", f"{h}.{escalated}.{s}"):
        assert all(r is None for r in _all(mutant)), mutant


def test_alg_none_forgery_beats_the_naive_verifier():
    """FORGERY A. No signature at all, and v0 hands back role=admin."""
    assert jwt.split(F_NONE)[2] == b""
    v0, v1, v2, v3 = _all(F_NONE)
    assert v0 == GREEDY and v0["role"] == "admin"
    assert (v1, v2, v3) == (None, None, None)


def test_alg_none_forgery_ignores_the_key_entirely():
    """The verification key is not an input to the decision (section 6.2)."""
    for key in (jwt.PUBLIC_KEY, b"", b"not a key at all", jwt.SECRET):
        assert jwt.verify_naive(F_NONE, key, NOW) == GREEDY, key


def test_key_confusion_beats_the_allowlist():
    """THE AHA. Signed with the RSA *public* key as an HMAC secret, so it
    needs no secret at all -- and v2's allowlist of two strong algorithms
    accepts it anyway.
    """
    assert len(jwt.split(F_CONF)[2]) == 32          # a real HMAC-SHA-256
    v0, v1, v2, v3 = _all(F_CONF)
    for r in (v0, v1, v2):
        assert r == GREEDY and r["role"] == "admin", r
    assert v3 is None


def test_the_load_bearing_line_is_where_alg_comes_from():
    """v2 and v3 differ in one line: header["alg"] vs. a caller constant.
    Both still allowlist, both still compare, both still check exp.
    """
    assert jwt.verify_allowlist(F_CONF, jwt.PUBLIC_KEY, NOW) is not None
    assert jwt.verify_strict(F_CONF, jwt.PUBLIC_KEY, "RS256", NOW) is None


def test_strict_must_pair_the_alg_with_the_key():
    """Section 7.2: naming an algorithm is only a fix if it is the one the
    key belongs to. "HS256" plus an RSA public key is the same bug.
    """
    assert jwt.verify_strict(F_CONF, jwt.PUBLIC_KEY, "HS256", NOW) == GREEDY
    assert jwt.verify_strict(F_CONF, jwt.SECRET, "HS256", NOW) is None


def test_key_confusion_needs_the_bytes_exactly():
    """Boundary 1 (section 6.7). Not "knows the public key" -- knows the
    exact serialisation the server feeds to HMAC.
    """
    base = jwt.PUBLIC_KEY
    for variant in (base + b"\n",
                    base.replace(b"3351", b"3352", 1),
                    json.dumps({"kty": "RSA", "e": jwt.E, "n": jwt.N},
                               separators=(",", ":")).encode(),
                    str(jwt.N).encode()):
        token = jwt.encode(GREEDY, "HS256", variant)
        assert jwt.verify_allowlist(token, base, NOW) is None, variant
    assert jwt.verify_allowlist(jwt.encode(GREEDY, "HS256", base),
                                base, NOW) == GREEDY


def test_symmetric_deployments_are_immune_to_key_confusion():
    """Boundary 2 (section 6.7). Nothing public to downgrade to -- but
    alg:none still walks in, which is why it needed its own fix.
    """
    for guess in (b"", b"secret", jwt.PUBLIC_KEY):
        token = jwt.encode(GREEDY, "HS256", guess)
        assert jwt.verify_allowlist(token, jwt.SECRET, NOW) is None, guess
    assert jwt.verify_naive(F_NONE, jwt.SECRET, NOW) == GREEDY
    assert jwt.verify_blocklist(F_NONE, jwt.SECRET, NOW) is None


def test_the_header_is_signed_data():
    """Relabelling a header breaks the signature: the signing input is the
    encoded text, not the decoded JSON (section 7.4).
    """
    _h, p, s = F_CONF.split(".")
    for header in ({"alg": "RS256", "typ": "JWT"}, {"alg": "HS256", "typ": "JWS"}):
        relabelled = f"{jwt.b64e(jwt._json(header))}.{p}.{s}"
        assert all(r is None for r in _all(relabelled)), header


def test_exp_is_only_an_expiry_if_it_is_required():
    """Section 7.5. `if "exp" in payload` is not an expiry check."""
    forever = jwt.encode({"sub": "alice", "role": "admin"}, "HS256", jwt.SECRET)
    far_future = NOW + 86400 * 3650
    assert jwt.verify_naive(forever, jwt.SECRET, far_future) is not None
    assert jwt.verify_strict(forever, jwt.SECRET, "HS256", far_future) is None
    expired = jwt.encode({"sub": "a", "exp": NOW - 1}, "HS256", jwt.SECRET)
    assert jwt.verify_naive(expired, jwt.SECRET, NOW) is None


def test_peek_needs_no_key():
    """base64 is not encryption (section 6.6)."""
    assert jwt.peek(RS_TOKEN) == ({"alg": "RS256", "typ": "JWT"}, HONEST)


def test_wrong_signatures_are_still_rejected():
    """The guard rail: none of this is "the verifiers accept everything"."""
    bad_hmac = jwt.encode(GREEDY, "HS256", b"a guess")
    bad_rsa = RS_TOKEN.rsplit(".", 1)[0] + "." + jwt.b64e(b"\x00" * 64)
    unknown = f"{jwt.b64e(jwt._json({'alg': 'HS512'}))}.{RS_TOKEN.split('.')[1]}."
    for token in (bad_hmac, bad_rsa, unknown):
        assert all(r is None for r in _all(token)), token


def test_rsa_round_trips_and_is_not_a_stub():
    """The toy RSA really is asymmetric: signing needs D, checking needs E."""
    si = b"header.payload"
    sig = jwt.rs256_sign(si, jwt.PRIVATE_KEY)
    assert jwt.rs256_check(si, sig, jwt.PUBLIC_KEY) is True
    assert jwt.rs256_check(b"header.payloae", sig, jwt.PUBLIC_KEY) is False
    assert jwt.N.bit_length() == 511 and jwt.P * jwt.Q == jwt.N


TESTS = [
    test_honest_tokens_round_trip,
    test_tampering_the_payload_is_caught,
    test_alg_none_forgery_beats_the_naive_verifier,
    test_alg_none_forgery_ignores_the_key_entirely,
    test_key_confusion_beats_the_allowlist,
    test_the_load_bearing_line_is_where_alg_comes_from,
    test_strict_must_pair_the_alg_with_the_key,
    test_key_confusion_needs_the_bytes_exactly,
    test_symmetric_deployments_are_immune_to_key_confusion,
    test_the_header_is_signed_data,
    test_exp_is_only_an_expiry_if_it_is_required,
    test_peek_needs_no_key,
    test_wrong_signatures_are_still_rejected,
    test_rsa_round_trips_and_is_not_a_stub,
]


if __name__ == "__main__":
    failed = 0
    for test in TESTS:
        try:
            test()
        except AssertionError as exc:
            failed += 1
            print(f"FAIL  {test.__name__}: {exc}")
        else:
            print(f"PASS  {test.__name__}")
    if failed:
        print(f"\n{failed} of {len(TESTS)} tests FAILED")
        raise SystemExit(1)
    print(f"\nAll {len(TESTS)} tests PASSED")