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

This is a DELIBERATELY VULNERABLE server for educational use ONLY.
Run it only inside the provided lab network. Do NOT expose it to
any network you do not fully control.

The server holds a secret AES-CBC ciphertext. It exposes one endpoint:

  POST /check  {"iv": "<hex>", "ct": "<hex>"}
    -> 200 {"valid": true}   if PKCS#7 padding is correct
    -> 200 {"valid": false}  if PKCS#7 padding is wrong

That one-bit oracle is the entire attack surface.
"""

import os
from flask import Flask, request, jsonify
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend

app = Flask(__name__)

# Fixed AES-128 key (known to the server, unknown to the attacker)
KEY = bytes.fromhex("0123456789abcdef0123456789abcdef")

# The target plaintext (PKCS#7-padded to a multiple of 16 bytes)
PLAINTEXT = b"TARGETPLAINTXT!!"  # exactly 16 bytes, pad = 0x10 * 16
IV = os.urandom(16)
CIPHERTEXT = None


def pkcs7_pad(data: bytes, block_size: int = 16) -> bytes:
    pad_len = block_size - (len(data) % block_size)
    return data + bytes([pad_len] * pad_len)


def pkcs7_unpad(data: bytes) -> bytes:
    """Raise ValueError if padding is invalid."""
    if not data:
        raise ValueError("empty")
    pad_len = data[-1]
    if pad_len == 0 or pad_len > 16:
        raise ValueError("bad pad byte")
    if data[-pad_len:] != bytes([pad_len] * pad_len):
        raise ValueError("bad padding")
    return data[:-pad_len]


def aes_cbc_encrypt(key: bytes, iv: bytes, plaintext: bytes) -> bytes:
    padded = pkcs7_pad(plaintext)
    cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
    enc = cipher.encryptor()
    return enc.update(padded) + enc.finalize()


def aes_cbc_decrypt_raw(key: bytes, iv: bytes, ciphertext: bytes) -> bytes:
    """Decrypt without padding removal."""
    cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
    dec = cipher.decryptor()
    return dec.update(ciphertext) + dec.finalize()


@app.before_request
def _init():
    global CIPHERTEXT, IV
    if CIPHERTEXT is None:
        IV = os.urandom(16)
        CIPHERTEXT = aes_cbc_encrypt(KEY, IV, PLAINTEXT)


@app.route("/target", methods=["GET"])
def get_target():
    """Return the IV and ciphertext that the attacker must decrypt."""
    return jsonify({"iv": IV.hex(), "ct": CIPHERTEXT.hex()})


@app.route("/check", methods=["POST"])
def check_padding():
    """
    The oracle: tells the attacker only whether the submitted
    (IV, ciphertext) pair decrypts to a validly-padded block.
    """
    data = request.get_json(force=True)
    try:
        iv = bytes.fromhex(data["iv"])
        ct = bytes.fromhex(data["ct"])
    except (KeyError, ValueError):
        return jsonify({"error": "send {iv: hex, ct: hex}"}), 400

    if len(iv) != 16 or len(ct) % 16 != 0:
        return jsonify({"valid": False})

    decrypted = aes_cbc_decrypt_raw(KEY, iv, ct)
    try:
        pkcs7_unpad(decrypted)
        return jsonify({"valid": True})
    except ValueError:
        return jsonify({"valid": False})


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5000)
