Cryptography & Cache-Oblivious Algorithms

import hashlib, math, random

# ── RSA KEY GENERATION (educational, use cryptography library in prod) ──
def extended_gcd(a, b):
    """Extended Euclidean: returns (gcd, x, y) with a*x + b*y = gcd."""
    if a == 0: return b, 0, 1
    g, x, y = extended_gcd(b % a, a)
    return g, y - (b//a)*x, x

def mod_inverse(e, phi):
    """Compute d = e^(-1) mod phi via extended Euclidean."""
    g, d, _ = extended_gcd(e % phi, phi)
    if g != 1: raise ValueError("No inverse — gcd != 1")
    return d % phi

def rsa_keygen(p, q, e=65537):
    """
    RSA key generation from two primes p, q.
    Public key: (N, e). Private key: (d, p, q).
    Security relies on difficulty of factoring N = p*q.
    """
    N   = p * q
    phi = (p - 1) * (q - 1)
    assert math.gcd(e, phi) == 1, "e must be coprime to phi"
    d   = mod_inverse(e, phi)
    return (N, e), (d, p, q)   # (public_key, private_key)

def rsa_encrypt(m, public_key):
    N, e = public_key
    return pow(m, e, N)         # c = m^e mod N

def rsa_decrypt(c, private_key):
    d, p, q = private_key
    N = p * q
    return pow(c, d, N)         # m = c^d mod N

# Example (use large primes in practice!)
pk, sk = rsa_keygen(61, 53)
m   = 42
c   = rsa_encrypt(m, pk)
m2  = rsa_decrypt(c, sk)
print(f"Original: {m}, Ciphertext: {c}, Decrypted: {m2}")  # 42


# ── MERKLE TREE ──────────────────────────────────────────────────────
class MerkleTree:
    """
    Merkle tree for integrity verification of n data blocks.
    Trusted storage: O(1) — just the root hash.
    Verification of block i: O(log n) — path to root.
    Update of block i: O(log n) — recompute path.
    """
    def __init__(self, blocks):
        """blocks: list of bytes objects."""
        self.n      = len(blocks)
        self.leaves = [hashlib.sha256(b).hexdigest() for b in blocks]
        self.tree   = self._build(self.leaves)

    def _hash_pair(self, left, right):
        return hashlib.sha256((left+right).encode()).hexdigest()

    def _build(self, layer):
        if len(layer) == 1: return layer
        if len(layer) % 2 == 1: layer = layer + [layer[-1]]  # duplicate last
        next_layer = [self._hash_pair(layer[i], layer[i+1])
                      for i in range(0, len(layer), 2)]
        return self._build(next_layer) + layer

    @property
    def root(self):
        return self.tree[0]  # root hash (store locally — O(1))

    def proof(self, i):
        """Generate Merkle proof: O(log n) sibling hashes."""
        idx   = len(self.tree) - self.n + i
        path  = []
        while idx > 0:
            sibling = idx - 1 if idx % 2 == 1 else idx + 1
            if sibling < len(self.tree):
                path.append((sibling, self.tree[sibling]))
            idx = (idx - 1) // 2
        return path

    def verify(self, block, i, proof, root):
        """Verify block i against stored root hash. O(log n)."""
        current = hashlib.sha256(block).hexdigest()
        idx     = len(self.tree) - self.n + i
        for sib_idx, sib_hash in proof:
            if sib_idx == idx - 1:  # sibling is left
                current = self._hash_pair(sib_hash, current)
            else:
                current = self._hash_pair(current, sib_hash)
            idx = (idx - 1) // 2
        return current == root

# Example: 4 files
files  = [b"file0", b"file1", b"file2", b"file3"]
tree   = MerkleTree(files)
print(f"Root hash: {tree.root[:16]}...")

proof  = tree.proof(1)
valid  = tree.verify(b"file1", 1, proof, tree.root)
print(f"Proof valid: {valid}")           # True

tampered = tree.verify(b"TAMPERED", 1, proof, tree.root)
print(f"Tampered detected: {not tampered}")  # True


# ── DIFFIE-HELLMAN KEY EXCHANGE ───────────────────────────────────────
def diffie_hellman(p, g):
    """
    Diffie-Hellman key exchange over Z_p^*.
    p: large prime, g: generator.
    Returns (alice_secret, bob_secret, shared_key).
    Security: discrete logarithm problem in Z_p^*.
    """
    a = random.randint(2, p-2)  # Alice's secret
    b = random.randint(2, p-2)  # Bob's secret
    A = pow(g, a, p)            # Alice sends g^a mod p
    B = pow(g, b, p)            # Bob sends g^b mod p
    alice_key = pow(B, a, p)    # Alice computes (g^b)^a = g^ab mod p
    bob_key   = pow(A, b, p)    # Bob computes (g^a)^b = g^ab mod p
    assert alice_key == bob_key
    return a, b, alice_key

p  = 23       # toy prime (use 2048-bit primes in practice)
g  = 5        # generator
_, _, k = diffie_hellman(p, g)
print(f"DH shared key: {k}")

import numpy as np

# ── CACHE-OBLIVIOUS MATRIX MULTIPLICATION ────────────────────────────
# Track memory transfers (simulated with block access counting)
class CacheSimulator:
    """Simplified cache simulator to count block transfers."""
    def __init__(self, B=8, M=64):
        self.B, self.M = B, M
        self.cache = {}     # block_id -> last_access_time
        self.time  = 0
        self.transfers = 0

    def access(self, addr):
        """Access a word at address; fetch block if not in cache."""
        block_id = addr // self.B
        if block_id not in self.cache:
            self.transfers += 1
            if len(self.cache) >= self.M // self.B:
                # LRU eviction
                lru = min(self.cache, key=self.cache.get)
                del self.cache[lru]
            self.cache[block_id] = self.time
        else:
            self.cache[block_id] = self.time
        self.time += 1


def co_matmul(X, Y, B=8, M=64):
    """
    Cache-oblivious matrix multiplication via recursive subdivision.
    No explicit blocking — the recursion naturally achieves cache-efficiency.
    MT(N) = O(N^3 / (B * sqrt(M))) — optimal.

    For demonstration: pure Python recursive; NumPy is much faster in practice.
    """
    n = len(X)
    if n <= 1:
        return [[X[0][0] * Y[0][0]]]

    # Divide matrices into four n/2 x n/2 quadrants
    h = n // 2
    A11, A12 = [r[:h] for r in X[:h]], [r[h:] for r in X[:h]]
    A21, A22 = [r[:h] for r in X[h:]], [r[h:] for r in X[h:]]
    B11, B12 = [r[:h] for r in Y[:h]], [r[h:] for r in Y[:h]]
    B21, B22 = [r[:h] for r in Y[h:]], [r[h:] for r in Y[h:]]

    def add(P, Q):
        return [[P[i][j]+Q[i][j] for j in range(len(P[0]))]
                for i in range(len(P))]

    # 8 recursive multiplications (like standard divide-and-conquer)
    C11 = add(co_matmul(A11,B11), co_matmul(A12,B21))
    C12 = add(co_matmul(A11,B12), co_matmul(A12,B22))
    C21 = add(co_matmul(A21,B11), co_matmul(A22,B21))
    C22 = add(co_matmul(A21,B12), co_matmul(A22,B22))

    # Combine quadrants
    result = [[0]*n for _ in range(n)]
    for i in range(h):
        result[i][:h]   = C11[i]
        result[i][h:]   = C12[i]
        result[h+i][:h] = C21[i]
        result[h+i][h:] = C22[i]
    return result


# ── VAN EMDE BOAS TREE LAYOUT ─────────────────────────────────────────
def veb_layout(sorted_arr):
    """
    Lay out a sorted array in van Emde Boas (vEB) tree order.
    This gives O(log_B N) memory transfers for binary search,
    without knowing B — cache-oblivious.

    Recursively: take middle element as root, recurse on left half
    (bottom subtrees) and right half (top subtrees).
    Each subtree of size B fits in O(1) cache blocks.
    """
    n = len(sorted_arr)
    if n <= 1:
        return sorted_arr
    if n == 2:
        return [sorted_arr[0], sorted_arr[1]]

    # Split at the midpoint of tree height (not array midpoint)
    import math
    h     = int(math.log2(n + 1))    # height of BST
    h_top = h // 2
    h_bot = h - h_top
    top_size = (1 << h_top) - 1      # nodes in top subtree
    bot_size = (1 << h_bot) - 1      # nodes in each bottom subtree

    # Root of top subtree is at position top_size in a BFS ordering
    # vEB layout recursively arranges bottom subtrees, then top subtree
    top_subtree = sorted_arr[:top_size + 1]
    bot_subtrees = [sorted_arr[top_size + 1 + i*(bot_size+1):
                               top_size + 1 + (i+1)*(bot_size+1)]
                    for i in range(top_size + 1)]

    result = []
    for sub in bot_subtrees:
        result.extend(veb_layout(sub))
    result.extend(veb_layout(top_subtree))
    return result


# ── BINARY MERGESORT TRANSFER ANALYSIS ───────────────────────────────
def mergesort_transfers(N, B, M):
    """
    Estimate memory transfers for binary mergesort.
    MT(N) = O(N/B * log_{M/B}(N/M)) — not optimal.
    Optimal: N/B * log_{M/B}(N/B).
    """
    if N <= M:
        return N / B    # fits in cache: O(N/B) to load
    # log_{M/B}(N/M) levels, each costs N/B
    levels = math.log(N/M, M/B)
    return N / B * levels

import math
N, B, M = 1_000_000, 64, 8192
print(f"Binary mergesort transfers: {mergesort_transfers(N,B,M):.0f}")
opt = N/B * math.log(N/B, M/B)
print(f"Optimal: {opt:.0f}")
print(f"Ratio: {mergesort_transfers(N,B,M)/opt:.2f}")

import hmac, hashlib, json, time

# ── HMAC FOR TRADE MESSAGE AUTHENTICATION ────────────────────────────
def sign_order(order_dict, secret_key: bytes) -> str:
    """
    HMAC-SHA256 authentication of a trade order message.
    Used in XTS API, FIX protocol, and most broker APIs.

    Properties:
    - Correctness: same key always produces same MAC
    - Unforgeability: without secret_key, cannot produce valid MAC
    - Unlike plain hash: adversary who knows MAC(k, m1) cannot
      compute MAC(k, m2) for any new m2 (security from secret key)
    """
    message = json.dumps(order_dict, sort_keys=True).encode()
    mac = hmac.new(secret_key, message, hashlib.sha256).hexdigest()
    return mac

def verify_order(order_dict, received_mac: str, secret_key: bytes) -> bool:
    """
    Verify HMAC. Use hmac.compare_digest to prevent timing attacks.
    Timing attack: if we return early on first mismatch, an attacker
    can deduce the correct MAC one byte at a time by measuring response time.
    """
    expected = sign_order(order_dict, secret_key)
    return hmac.compare_digest(expected, received_mac)  # constant-time compare!


# Example: XTS-style order authentication
secret_key = b'ttt-secret-2025-xts-api-key'

order = {
    'order_id':   'TTT-20250327-001',
    'instrument': 'NIFTY25MAR24000CE',
    'quantity':   50,
    'price':      245.50,
    'side':       'BUY',
    'timestamp':  int(time.time()),
}
mac = sign_order(order, secret_key)
print(f"Order MAC (first 16 chars): {mac[:16]}...")

# Verify legitimate order
print(f"Legitimate order valid: {verify_order(order, mac, secret_key)}")

# Tampered order (e.g., price changed by attacker)
tampered_order = dict(order, price=1.0)   # attacker changes price
print(f"Tampered order valid: {verify_order(tampered_order, mac, secret_key)}")
# False — HMAC detects tampering


# ── COMMITMENT SCHEME ─────────────────────────────────────────────────
import os
def commit(value: str) -> tuple:
    """
    Cryptographic commitment: hide value v until reveal.
    Used in: sealed-bid auctions, fair coin flips, zero-knowledge proofs.
    commit(v) = (commitment=h(r||v), randomness=r) — HIDING: h(r||v) leaks nothing
    reveal(v, r): check commitment = h(r||v) — BINDING: can't find v' with h(r||v')=h(r||v)
    """
    r = os.urandom(32)   # 256 bits of randomness (prevents brute force)
    commitment = hashlib.sha256(r + value.encode()).hexdigest()
    return commitment, r

def reveal(value: str, randomness: bytes, commitment: str) -> bool:
    """Verify that commitment was to value."""
    expected = hashlib.sha256(randomness + value.encode()).hexdigest()
    return hmac.compare_digest(expected, commitment)

# TTT use case: sealed-bid order submission
bid_price = "24500"
c, r = commit(bid_price)
print(f"Committed to bid. Commitment: {c[:16]}...")
# ... later, when all bids are submitted ...
print(f"Reveal valid: {reveal(bid_price, r, c)}")       # True
print(f"Cheat attempt: {reveal('23000', r, c)}")        # False