import base64
import json

from mitmproxy import http
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes

KEY = b"0123456789abcdef"
IV = b"fedcba9876543210"


def encrypt(text: str) -> str:
    data = text.encode()
    n = 16 - (len(data) % 16)
    data += bytes([n] * n)
    enc = Cipher(algorithms.AES(KEY), modes.CBC(IV)).encryptor()
    return base64.b64encode(enc.update(data) + enc.finalize()).decode()


def decrypt(b64: str) -> str:
    dec = Cipher(algorithms.AES(KEY), modes.CBC(IV)).decryptor()
    pt = dec.update(base64.b64decode(b64)) + dec.finalize()
    return pt[: -pt[-1]].decode()


def request(flow: http.HTTPFlow) -> None:
    if flow.request.headers.get("X-Encrypted") != "1":
        return
    try:
        envelope = json.dumps({"payload": encrypt(flow.request.get_text())})
        flow.request.set_text(envelope)                  # server gets ciphertext
        print("[server-side] request re-encrypted for server")
    except Exception as exc:
        print(f"[server-side] request error: {exc}")


def response(flow: http.HTTPFlow) -> None:
    if flow.request.headers.get("X-Encrypted") != "1":
        return
    try:
        payload = json.loads(flow.response.get_text())["payload"]
        flow.response.set_text(decrypt(payload))         # Burp sees cleartext
        print("[server-side] response decrypted for Burp")
    except Exception as exc:
        print(f"[server-side] response error: {exc}")