Cryptographic Implementation Security Questions

import hashlib
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.ciphers.aead import AESGCM

KEY = hashlib.sha256(b"pinned-demo-key-seed-v2").digest()


def ecb_encrypt(plaintext_blocks: bytes) -> bytes:
    encryptor = Cipher(algorithms.AES(KEY), modes.ECB()).encryptor()
    return encryptor.update(plaintext_blocks) + encryptor.finalize()


def aead_encrypt(plaintext: bytes, nonce: bytes) -> bytes:
    return AESGCM(KEY).encrypt(nonce, plaintext, associated_data=None)


if __name__ == "__main__":
    # --- misuse 1: ECB mode leaks that two 16-byte blocks are equal, even
    # though the attacker never learns the key or the plaintext values ---
    block = b"REPEAT-PATTERN!!"  # 16 bytes, deliberately identical across repeats
    plaintext = block * 4        # 4 identical plaintext blocks, e.g. a bitmap's blank rows
    ct_ecb = ecb_encrypt(plaintext)
    ciphertext_blocks = [ct_ecb[i:i + 16] for i in range(0, len(ct_ecb), 16)]
    print("ECB ciphertext blocks for 4 identical plaintext blocks:")
    for i, b in enumerate(ciphertext_blocks):
        print(f"  block {i}: {b.hex()}")
    print(f"all 4 ciphertext blocks identical: {len(set(ciphertext_blocks)) == 1} "
          f"(structure of the plaintext leaks with zero key knowledge)")

    # --- fix: AEAD (AES-GCM). Same repeating plaintext, unique nonce, and the
    # ciphertext no longer exposes block-level repetition ---
    nonce = hashlib.sha256(b"pinned-demo-nonce-seed-v2").digest()[:12]
    ct_gcm = aead_encrypt(plaintext, nonce)
    gcm_body = ct_gcm[:-16]  # strip the 16-byte auth tag to compare ciphertext body only
    gcm_blocks = [gcm_body[i:i + 16] for i in range(0, len(gcm_body), 16)]
    print(f"\nAES-GCM ciphertext blocks for the SAME repeating plaintext:")
    for i, b in enumerate(gcm_blocks):
        print(f"  block {i}: {b.hex()}")
    print(f"all 4 ciphertext blocks identical: {len(set(gcm_blocks)) == 1} (must be False)")

    # --- misuse 2: nonce reuse under a stream-like mode (CTR) leaks the XOR
    # of two plaintexts the moment the SAME nonce is reused ---
    def ctr_encrypt(pt: bytes, nonce: bytes) -> bytes:
        encryptor = Cipher(algorithms.AES(KEY), modes.CTR(nonce + b"\x00" * (16 - len(nonce)))).encryptor()
        return encryptor.update(pt) + encryptor.finalize()

    reused_nonce = hashlib.sha256(b"pinned-demo-nonce-seed-v2").digest()[:12]
    msg_a = b"transfer $10 to bob!!!"
    msg_b = b"transfer $99999 to eve"
    assert len(msg_a) == len(msg_b)

    ct_a = ctr_encrypt(msg_a, reused_nonce)
    ct_b = ctr_encrypt(msg_b, reused_nonce)  # same nonce reused: the bug

    xor_of_ciphertexts = bytes(a ^ b for a, b in zip(ct_a, ct_b))
    xor_of_plaintexts = bytes(a ^ b for a, b in zip(msg_a, msg_b))
    print(f"\nnonce reuse under CTR: XOR(ct_a, ct_b) == XOR(pt_a, pt_b): "
          f"{xor_of_ciphertexts == xor_of_plaintexts}")
    print("(an attacker who knows or guesses one plaintext recovers the other "
          "directly from that XOR, with no key material needed)")

ECB ciphertext blocks for 4 identical plaintext blocks:
  block 0: d7e3d212a2dc0b5cc4d3de6846be077d
  block 1: d7e3d212a2dc0b5cc4d3de6846be077d
  block 2: d7e3d212a2dc0b5cc4d3de6846be077d
  block 3: d7e3d212a2dc0b5cc4d3de6846be077d
all 4 ciphertext blocks identical: True (structure of the plaintext leaks with zero key knowledge)

AES-GCM ciphertext blocks for the SAME repeating plaintext:
  block 0: 544f53cb313ae00cf273aa7c5c20ef99
  block 1: 05a2493f93909dcb04083f6b8781d9b0
  block 2: 7b0700634e38a35c2fd4bf9de50fdee8
  block 3: 5650b8a9e43732616eb5f999e4a39f4e
all 4 ciphertext blocks identical: False (must be False)

nonce reuse under CTR: XOR(ct_a, ct_b) == XOR(pt_a, pt_b): True
(an attacker who knows or guesses one plaintext recovers the other directly from that XOR, with no key material needed)

import hmac
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives import padding as sympadding

KEY = bytes(range(32))          # AES-256 key, fixed for reproducibility (demo only)
MAC_KEY = bytes(range(32, 64))  # separate HMAC key -- never reuse the encryption key
IV = bytes(range(16))

def aes_cbc_decrypt(ciphertext, key, iv):
    d = Cipher(algorithms.AES(key), modes.CBC(iv)).decryptor()
    return d.update(ciphertext) + d.finalize()

def aes_cbc_encrypt(padded, key, iv):
    e = Cipher(algorithms.AES(key), modes.CBC(iv)).encryptor()
    return e.update(padded) + e.finalize()

def pkcs7_unpad(padded):
    u = sympadding.PKCS7(128).unpadder()
    # .finalize() performs the actual padding check -- omitting it silently accepts
    # any trailing bytes and would defeat this whole example.
    return u.update(padded) + u.finalize()

def pkcs7_pad(data):
    p = sympadding.PKCS7(128).padder()
    return p.update(data) + p.finalize()

def hmac_tag(data, mac_key):
    return hmac.new(mac_key, data, digestmod='sha256').digest()

def vulnerable_decrypt(ciphertext, tag, key, iv, mac_key):
    """As given in the question: decrypt, unpad, THEN check the MAC, with a distinct
    return string for each failure reason."""
    plaintext_padded = aes_cbc_decrypt(ciphertext, key, iv)
    try:
        unpadded = pkcs7_unpad(plaintext_padded)
    except ValueError:
        return 'padding error'
    expected = hmac_tag(ciphertext, mac_key)
    if not hmac.compare_digest(expected, tag):
        return 'mac error'
    return unpadded

def fixed_decrypt(ciphertext, tag, key, iv, mac_key):
    """Verify-then-decrypt: the MAC is checked over the ciphertext, in constant time,
    before the decryptor or unpadder ever runs. Any failure -- bad tag or (only
    theoretically reachable) bad padding under a valid tag -- returns the same None."""
    expected = hmac_tag(ciphertext, mac_key)
    if not hmac.compare_digest(expected, tag):
        return None
    plaintext_padded = aes_cbc_decrypt(ciphertext, key, iv)
    try:
        return pkcs7_unpad(plaintext_padded)
    except ValueError:
        return None

def run_case(label, ciphertext, tag):
    v = vulnerable_decrypt(ciphertext, tag, KEY, IV, MAC_KEY)
    f = fixed_decrypt(ciphertext, tag, KEY, IV, MAC_KEY)
    print(f"{label}:")
    print(f"  vulnerable_decrypt -> {v!r}")
    print(f"  fixed_decrypt      -> {f!r}")

plaintext = b'transfer:5000:acct-8842-checking'
padded = pkcs7_pad(plaintext)
ct = aes_cbc_encrypt(padded, KEY, IV)
good_tag = hmac_tag(ct, MAC_KEY)

run_case('A. valid ciphertext + valid tag', ct, good_tag)

early_flip = bytearray(ct)
early_flip[0] ^= 0x01  # perturbs only the first plaintext block, padding intact
run_case('B. bit-flip in an early block (padding untouched, MAC invalid)',
          bytes(early_flip), good_tag)

bad_padded = bytearray(padded)
bad_padded[-1] = 0xFF  # not a legal PKCS7 pad byte
bad_ct = aes_cbc_encrypt(bytes(bad_padded), KEY, IV)
bad_tag = hmac_tag(bad_ct, MAC_KEY)
run_case('C. valid tag, malformed padding underneath', bad_ct, bad_tag)

A. valid ciphertext + valid tag:
  vulnerable_decrypt -> b'transfer:5000:acct-8842-checking'
  fixed_decrypt      -> b'transfer:5000:acct-8842-checking'
B. bit-flip in an early block (padding untouched, MAC invalid):
  vulnerable_decrypt -> 'mac error'
  fixed_decrypt      -> None
C. valid tag, malformed padding underneath:
  vulnerable_decrypt -> 'padding error'
  fixed_decrypt      -> None

import secrets
import hmac
import hashlib
import time
import random


class ResetTokenStore:
    """Server-side store: never persists the raw token, only its hash, plus
    expiry and a used flag for single-use enforcement."""

    def __init__(self):
        self._records = {}  # token_hash -> {"user": ..., "expires_at": ..., "used": bool}

    @staticmethod
    def _hash(raw_token: str) -> str:
        return hashlib.sha256(raw_token.encode()).hexdigest()

    def issue(self, user_id: str, ttl_seconds: int = 900) -> str:
        raw_token = secrets.token_urlsafe(32)  # 256 bits of CSPRNG entropy, URL-safe
        self._records[self._hash(raw_token)] = {
            "user": user_id,
            "expires_at": time.time() + ttl_seconds,
            "used": False,
        }
        return raw_token

    def redeem(self, raw_token: str) -> str | None:
        token_hash = self._hash(raw_token)
        record = self._records.get(token_hash)
        if record is None:
            return None
        if record["used"] or time.time() > record["expires_at"]:
            return None
        record["used"] = True  # single-use enforcement
        return record["user"]

    def verify_constant_time(self, raw_token: str, claimed_hash: str) -> bool:
        """Illustrates comparing token hashes without a short-circuiting ==,
        which would leak how many leading bytes matched via timing."""
        return hmac.compare_digest(self._hash(raw_token), claimed_hash)


class WeakResetTokenGenerator:
    """The anti-pattern under test: a 6-digit numeric code seeded from a
    non-cryptographic PRNG, the kind of 'looks fine in a demo' shortcut this
    question is warning against."""

    def __init__(self, seed):
        self._rng = random.Random(seed)  # NOT a CSPRNG

    def generate(self) -> str:
        return f"{self._rng.randrange(0, 1_000_000):06d}"


if __name__ == "__main__":
    store = ResetTokenStore()

    # --- correct path ---
    token = store.issue("user-42", ttl_seconds=900)
    print(f"issued token length: {len(token)} chars, generated from 32 random "
          f"bytes = 256 bits of entropy (the specific characters are genuine "
          f"CSPRNG output and differ on every run)")

    user = store.redeem(token)
    print(f"first redeem succeeds, resolves to user: {user}")

    stored_hash = store._hash(token)
    print(f"constant-time hash comparison result for the correct token: "
          f"{store.verify_constant_time(token, stored_hash)}")
    print(f"constant-time hash comparison result for a wrong token: "
          f"{store.verify_constant_time('not-the-token', stored_hash)}")

    user_again = store.redeem(token)
    print(f"second redeem of the SAME token (replay): {user_again} (must be None, single-use holds)")

    # --- expiry enforcement ---
    short_lived_store = ResetTokenStore()
    short_token = short_lived_store.issue("user-99", ttl_seconds=0)
    time.sleep(0.01)
    expired_result = short_lived_store.redeem(short_token)
    print(f"redeem after expiry: {expired_result} (must be None)")

    # --- attacking case: brute-force the weak 6-digit generator's keyspace
    # (1,000,000 values), a search size that is realistically enumerable by
    # an attacker hammering the reset endpoint if it isn't rate-limited ---
    weak_gen = WeakResetTokenGenerator(seed=1)
    real_code = weak_gen.generate()
    print(f"\nweak generator's issued code: {real_code}")

    guesses_tried = 0
    found = False
    for candidate in range(1_000_000):
        guesses_tried += 1
        if f"{candidate:06d}" == real_code:
            found = True
            break
    print(f"brute-forced the 6-digit code in {guesses_tried} guesses "
          f"(worst case {1_000_000}; this is why rate-limiting AND a large "
          f"keyspace are both required, and why the numeric-code shortcut is unsafe "
          f"without one)")

    # --- contrast: the CSPRNG token's keyspace is 2^256; demonstrate the
    # generator produces no repeats across a large honest sample (a collision
    # would falsify the "effectively unique" claim; we do not claim the full
    # keyspace is unsearchable, only report the observed collision count) ---
    rng_check = [secrets.token_urlsafe(32) for _ in range(50_000)]
    print(f"\ncollisions observed across {len(rng_check)} CSPRNG-generated tokens: "
          f"{len(rng_check) - len(set(rng_check))}")

issued token length: 43 chars, generated from 32 random bytes = 256 bits of entropy (the specific characters are genuine CSPRNG output and differ on every run)
first redeem succeeds, resolves to user: user-42
constant-time hash comparison result for the correct token: True
constant-time hash comparison result for a wrong token: False
second redeem of the SAME token (replay): None (must be None, single-use holds)
redeem after expiry: None (must be None)

weak generator's issued code: 140891
brute-forced the 6-digit code in 140892 guesses (worst case 1000000; this is
why rate-limiting AND a large keyspace are both required, and why the
numeric-code shortcut is unsafe without one)

collisions observed across 50000 CSPRNG-generated tokens: 0

import base64, json, hmac, hashlib
import jwt
from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.exceptions import InvalidSignature

# ---- pinned RSA keypair (2048-bit, real keygen) ----
private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
public_key = private_key.public_key()
priv_pem = private_key.private_bytes(
    encoding=serialization.Encoding.PEM,
    format=serialization.PrivateFormat.TraditionalOpenSSL,
    encryption_algorithm=serialization.NoEncryption(),
)
pub_pem = public_key.public_bytes(
    encoding=serialization.Encoding.PEM,
    format=serialization.PublicFormat.SubjectPublicKeyInfo,
)

payload = {"sub": "user-42", "role": "user"}
legit_token = jwt.encode(payload, priv_pem, algorithm="RS256")
print("legit RS256 token (truncated):", legit_token[:40], "...")


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


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


# ---- VULNERABLE verifier: dispatches on the attacker-supplied `alg` header,
# reusing one "key" for whichever code path the header names ----
def verify_vulnerable(token: str, key_material: bytes):
    header_b64, payload_b64, sig_b64 = token.split(".")
    header = json.loads(b64url_decode(header_b64))
    signing_input = f"{header_b64}.{payload_b64}".encode()
    alg = header["alg"]  # attacker-controlled

    if alg == "none":
        pass  # bug: "none" is honored, no signature check at all
    elif alg == "HS256":
        expected = hmac.new(key_material, signing_input, hashlib.sha256).digest()
        if not hmac.compare_digest(expected, b64url_decode(sig_b64)):
            raise ValueError("bad HMAC signature")
    elif alg == "RS256":
        pub = serialization.load_pem_public_key(key_material)
        pub.verify(b64url_decode(sig_b64), signing_input, padding.PKCS1v15(), hashes.SHA256())
    else:
        raise ValueError(f"unsupported alg {alg}")

    return json.loads(b64url_decode(payload_b64))


# ---- FIXED verifier: server pins the ONE expected algorithm, ignores the header ----
def verify_fixed(token: str, rsa_public_key_pem: bytes):
    return jwt.decode(token, rsa_public_key_pem, algorithms=["RS256"])


# --- Attack 1: alg=none forgery -------------------------------------------------
forged_header = b64url(json.dumps({"alg": "none", "typ": "JWT"}).encode())
forged_payload = b64url(json.dumps({"sub": "user-42", "role": "admin"}).encode())
forged_none_token = f"{forged_header}.{forged_payload}."  # empty signature segment

print("\n--- Attack 1: alg=none ---")
try:
    forged_claims = verify_vulnerable(forged_none_token, pub_pem)
    print("VULNERABLE verifier accepted forged token. Claims:", forged_claims)
except Exception as e:
    print("VULNERABLE verifier rejected it:", e)

try:
    verify_fixed(forged_none_token, pub_pem)
    print("FIXED verifier: WRONGLY accepted (should not happen)")
except Exception as e:
    print("FIXED verifier correctly rejected alg=none forgery:", type(e).__name__)

# --- Attack 2: RS256 -> HS256 algorithm confusion -------------------------------
# Attacker signs a token with HS256 using the SERVER'S PUBLIC KEY BYTES as the
# HMAC secret (public, so the attacker has it). verify_vulnerable's HS256
# branch reuses that same key material as an HMAC secret and will match.
forged_header2 = b64url(json.dumps({"alg": "HS256", "typ": "JWT"}).encode())
forged_payload2 = b64url(json.dumps({"sub": "user-42", "role": "admin"}).encode())
signing_input2 = f"{forged_header2}.{forged_payload2}".encode()
forged_sig2 = hmac.new(pub_pem, signing_input2, hashlib.sha256).digest()
forged_hs256_token = f"{forged_header2}.{forged_payload2}.{b64url(forged_sig2)}"

print("\n--- Attack 2: RS256->HS256 confusion ---")
try:
    forged_claims = verify_vulnerable(forged_hs256_token, pub_pem)
    print("VULNERABLE verifier accepted forged token. Claims:", forged_claims)
except Exception as e:
    print("VULNERABLE verifier rejected it:", e)

try:
    verify_fixed(forged_hs256_token, pub_pem)
    print("FIXED verifier: WRONGLY accepted (should not happen)")
except Exception as e:
    print("FIXED verifier correctly rejected HS256-confusion forgery:", type(e).__name__)

# --- Control: the fixed verifier still accepts the legitimate token ------------
print("\n--- Control: legitimate token against FIXED verifier ---")
claims = verify_fixed(legit_token, pub_pem)
print("FIXED verifier accepted the real token. Claims:", claims)

legit RS256 token (truncated): eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJ ...

--- Attack 1: alg=none ---
VULNERABLE verifier accepted forged token. Claims: {'sub': 'user-42', 'role': 'admin'}
FIXED verifier correctly rejected alg=none forgery: InvalidAlgorithmError

--- Attack 2: RS256->HS256 confusion ---
VULNERABLE verifier accepted forged token. Claims: {'sub': 'user-42', 'role': 'admin'}
FIXED verifier correctly rejected HS256-confusion forgery: InvalidAlgorithmError

--- Control: legitimate token against FIXED verifier ---
FIXED verifier accepted the real token. Claims: {'sub': 'user-42', 'role': 'user'}

import hmac
import hashlib
import struct


def naive_sign(key: bytes, purpose: str, payload: str) -> bytes:
    """VULNERABLE: purpose and payload are just concatenated. Two different
    (purpose, payload) pairs can produce the IDENTICAL signed bytestring."""
    message = (purpose + payload).encode()
    return hmac.new(key, message, hashlib.sha256).digest()


def domain_separated_sign(key: bytes, purpose: str, payload: str) -> bytes:
    """FIXED: length-prefix each field so no byte sequence is ambiguous
    between two different (purpose, payload) splits, AND derive a
    purpose-scoped subkey via HKDF so a session-token signature and a
    reset-link signature are never even computed under the same key."""
    subkey = hashlib.pbkdf2_hmac("sha256", key, purpose.encode(), 1)  # stand-in for HKDF-Expand(key, purpose)
    framed = struct.pack(">I", len(purpose)) + purpose.encode() + struct.pack(">I", len(payload)) + payload.encode()
    return hmac.new(subkey, framed, hashlib.sha256).digest()


if __name__ == "__main__":
    shared_key = b"user-42-per-user-secret-key-material"

    # --- the confusion: two DIFFERENT logical messages collide because the
    # naive concatenation has no boundary between purpose and payload ---
    session_sig = naive_sign(shared_key, "session:", "adminuser")
    reset_sig = naive_sign(shared_key, "session:a", "dminuser")
    print("naive concatenation, two different (purpose, payload) pairs:")
    print(f"  sign('session:', 'adminuser')  = {session_sig.hex()[:16]}...")
    print(f"  sign('session:a', 'dminuser')  = {reset_sig.hex()[:16]}...")
    print(f"  signatures identical: {session_sig == reset_sig}")

    # --- the concrete attack this enables: attacker holds a validly-issued
    # reset-link signature for a payload they control, and re-presents it as
    # a forged session token, if the verifier only checks the HMAC without
    # confirming which purpose issued it (the missing check is exactly what
    # a shared key make impossible to add cheaply) ---
    attacker_controlled_reset_payload = "a:admin-session-hijack"
    forged_reset_sig = naive_sign(shared_key, "reset:", attacker_controlled_reset_payload)
    equivalent_session_sig = naive_sign(shared_key, "reset:a", ":admin-session-hijack")
    print(f"\nattacker-obtained reset signature reused as a session signature: "
          f"{forged_reset_sig == equivalent_session_sig}")

    # --- the fix: domain-separated, purpose-scoped signing ---
    fixed_session_sig = domain_separated_sign(shared_key, "session:", "adminuser")
    fixed_reset_sig = domain_separated_sign(shared_key, "session:a", "dminuser")
    print(f"\ndomain-separated construction, same ambiguous split as above:")
    print(f"  sign('session:', 'adminuser')  = {fixed_session_sig.hex()[:16]}...")
    print(f"  sign('session:a', 'dminuser')  = {fixed_reset_sig.hex()[:16]}...")
    print(f"  signatures identical: {fixed_session_sig == fixed_reset_sig} (must be False)")

    # --- and purpose-scoped subkeys mean even IDENTICAL payloads signed for
    # different purposes produce different tags, so a reset-link signature
    # can never verify as a session-token signature at all ---
    same_payload_session = domain_separated_sign(shared_key, "session", "adminuser")
    same_payload_reset = domain_separated_sign(shared_key, "reset", "adminuser")
    print(f"\nsame payload 'adminuser', purposes 'session' vs 'reset': "
          f"signatures identical: {same_payload_session == same_payload_reset} (must be False)")

naive concatenation, two different (purpose, payload) pairs:
  sign('session:', 'adminuser')  = 48280526e93ea281...
  sign('session:a', 'dminuser')  = 48280526e93ea281...
  signatures identical: True

attacker-obtained reset signature reused as a session signature: True

domain-separated construction, same ambiguous split as above:
  sign('session:', 'adminuser')  = 4d703ee8c96f4310...
  sign('session:a', 'dminuser')  = 11020fd1286362cb...
  signatures identical: False (must be False)

same payload 'adminuser', purposes 'session' vs 'reset': signatures identical: False (must be False)