"""
Lab 4.1 - CBC Padding Oracle Server
Chapter 4: Cryptography and Encryption: Foundations

This server deliberately implements the padding-oracle vulnerability:
a decryption endpoint that returns a DISTINGUISHABLE error for bad padding
versus bad content. This is the exact flaw Vaudenay described in 2002.

The secret was encrypted at startup; the key is never exposed.
The /encrypt endpoint gives you the ciphertext to attack.
The /check_padding endpoint is the oracle - it leaks one bit per query.

SAFETY NOTE: This is a training target. Run it only inside the lab network
(docker compose), never on a public interface.
"""

import os
import json
import base64
from flask import Flask, request, jsonify
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives import padding as sym_padding

app = Flask(__name__)

# The key is generated at container start and never logged or exported.
# The student "forgets" it immediately - they attack the oracle, not the key.
_KEY = os.urandom(16)  # AES-128

# The secret message to recover via the padding-oracle attack
SECRET_PLAINTEXT = b"ThePaddingOracleAttackRecoversThisWithoutTheKey!"

def pkcs7_pad(data: bytes, block_size: int = 16) -> bytes:
    padder = sym_padding.PKCS7(block_size * 8).padder()
    return padder.update(data) + padder.finalize()

def pkcs7_unpad(data: bytes, block_size: int = 16) -> bytes:
    unpadder = sym_padding.PKCS7(block_size * 8).unpadder()
    return unpadder.update(data) + unpadder.finalize()

def aes_cbc_encrypt(key: bytes, plaintext: bytes) -> tuple[bytes, bytes]:
    iv = os.urandom(16)
    padded = pkcs7_pad(plaintext)
    cipher = Cipher(algorithms.AES(key), modes.CBC(iv))
    encryptor = cipher.encryptor()
    ciphertext = encryptor.update(padded) + encryptor.finalize()
    return iv, ciphertext

def aes_cbc_decrypt_raw(key: bytes, iv: bytes, ciphertext: bytes) -> bytes:
    """Decrypt without removing padding - used internally by the oracle."""
    cipher = Cipher(algorithms.AES(key), modes.CBC(iv))
    decryptor = cipher.decryptor()
    return decryptor.update(ciphertext) + decryptor.finalize()


# Pre-encrypt the secret at startup
_IV, _CIPHERTEXT = aes_cbc_encrypt(_KEY, SECRET_PLAINTEXT)


@app.route("/encrypt", methods=["GET"])
def get_ciphertext():
    """
    Hand the attacker the IV and ciphertext they need to attack.
    The key is NEVER returned.
    """
    return jsonify({
        "iv": base64.b64encode(_IV).decode(),
        "ciphertext": base64.b64encode(_CIPHERTEXT).decode(),
        "block_size": 16,
        "note": "The key is not here. Use /check_padding to run the oracle attack."
    })


@app.route("/check_padding", methods=["POST"])
def check_padding():
    """
    THE ORACLE: takes an IV + ciphertext, attempts CBC decryption,
    and returns a DISTINGUISHABLE response for:
      - valid PKCS7 padding  -> {"padding": "valid"}
      - invalid PKCS7 padding -> {"padding": "invalid"}

    This is the vulnerability. A real server must NEVER distinguish these.
    The fix (demonstrated in Lab 4.1 step 5) is to use AES-GCM: the
    integrity check fails before any padding is ever examined.
    """
    body = request.get_json(force=True)
    try:
        iv = base64.b64decode(body["iv"])
        ciphertext = base64.b64decode(body["ciphertext"])
    except Exception:
        return jsonify({"error": "bad input"}), 400

    if len(iv) != 16 or len(ciphertext) % 16 != 0:
        return jsonify({"error": "bad length"}), 400

    # Decrypt to raw bytes (no padding removal yet)
    raw = aes_cbc_decrypt_raw(_KEY, iv, ciphertext)

    # Deliberately distinguishable padding check - THIS IS THE VULNERABILITY
    try:
        pkcs7_unpad(raw)
        return jsonify({"padding": "valid"})
    except Exception:
        return jsonify({"padding": "invalid"})


@app.route("/verify", methods=["POST"])
def verify_recovery():
    """
    After running the attack, submit your recovered plaintext here
    to confirm it matches the original secret.
    """
    body = request.get_json(force=True)
    candidate = body.get("plaintext", "")
    if isinstance(candidate, str):
        candidate = candidate.encode()
    if candidate == SECRET_PLAINTEXT:
        return jsonify({"result": "CORRECT - padding oracle attack succeeded!"})
    else:
        return jsonify({
            "result": "not quite",
            "hint": f"expected {len(SECRET_PLAINTEXT)} bytes"
        })


@app.route("/health", methods=["GET"])
def health():
    return jsonify({"status": "ok", "lab": "4.1-padding-oracle"})


if __name__ == "__main__":
    print("[oracle] AES-CBC padding-oracle server starting.")
    print("[oracle] Key is generated in memory and never logged.")
    print("[oracle] GET /encrypt for the ciphertext, POST /check_padding to query the oracle.")
    app.run(host="0.0.0.0", port=5000, debug=False)
