"""
Criptografia do payload enviado pra API do cliente.

Envelope (hibrida): AES-256-GCM nos dados, RSA-OAEP na chave AES.

Por que nao RSA puro: RSA-2048 com OAEP-SHA256 cifra no maximo 190 bytes.
Um payload de compras passa disso facil. O padrao — mesmo do TLS, PGP e JWE —
e sortear uma chave simetrica por mensagem, cifrar os dados com ela e mandar
essa chave cifrada com a publica do destinatario.

So o dono da chave privada abre. Nos ciframos e nao conseguimos desfazer,
que era o objetivo.
"""

import base64
import json
import os

from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import padding, rsa
from cryptography.hazmat.primitives.ciphers.aead import AESGCM

ALGORITMO = "RSA-OAEP-256+A256GCM"

_OAEP = padding.OAEP(
    mgf=padding.MGF1(algorithm=hashes.SHA256()),
    algorithm=hashes.SHA256(),
    label=None,
)


def carregar_chave_publica(pem: str):
    """Le a chave publica do cliente (formato PEM, vinda do configs.json)."""
    return serialization.load_pem_public_key(pem.encode())


def carregar_chave_privada(pem: str, senha: bytes | None = None):
    """So usada em teste. Em producao a privada fica com o cliente."""
    return serialization.load_pem_private_key(pem.encode(), password=senha)


def criptografar(dados: dict, chave_publica) -> dict:
    """
    Cifra o dict e devolve o envelope pronto pra mandar no corpo do POST.

    Sai assim:
        {
          "alg":   "RSA-OAEP-256+A256GCM",
          "chave": "<chave AES cifrada com a publica, base64>",
          "nonce": "<nonce do GCM, base64>",
          "dados": "<json cifrado + tag de autenticacao, base64>"
        }
    """
    chave_aes = AESGCM.generate_key(bit_length=256)
    nonce = os.urandom(12)
    corpo = json.dumps(dados, ensure_ascii=False, separators=(",", ":")).encode()

    cifrado = AESGCM(chave_aes).encrypt(nonce, corpo, None)
    chave_cifrada = chave_publica.encrypt(chave_aes, _OAEP)

    return {
        "alg": ALGORITMO,
        "chave": base64.b64encode(chave_cifrada).decode(),
        "nonce": base64.b64encode(nonce).decode(),
        "dados": base64.b64encode(cifrado).decode(),
    }


def descriptografar(envelope: dict, chave_privada) -> dict:
    """
    O lado do cliente. Fica aqui pra testar o round-trip e pra servir de
    referencia de implementacao pra quem for escrever a API que recebe.
    """
    chave_aes = chave_privada.decrypt(base64.b64decode(envelope["chave"]), _OAEP)
    corpo = AESGCM(chave_aes).decrypt(
        base64.b64decode(envelope["nonce"]),
        base64.b64decode(envelope["dados"]),
        None,
    )
    return json.loads(corpo)


def gerar_par_de_chaves() -> tuple[str, str]:
    """Gera um par RSA-2048 em PEM. Util pra homologacao: (privada, publica)."""
    chave = rsa.generate_private_key(public_exponent=65537, key_size=2048)
    privada = chave.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=serialization.NoEncryption(),
    ).decode()
    publica = (
        chave.public_key()
        .public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo,
        )
        .decode()
    )
    return privada, publica