#!/usr/bin/env python3
"""
Deszyfrowanie raportu retrospektywnego (RSA/AESG/SHA256).

Wymagania:
    pip install cryptography

Użycie:
    # 1. Generowanie pary kluczy (jednorazowo)
    openssl genrsa -out rsr-data.key 4096
    openssl rsa -in rsr-data.key -pubout -out rsr-data.pub.pem

    # 2. Zlecenie raportu (zapisuje zaszyfrowany klucz AES z odpowiedzi)
    curl -s -X POST https://<host>/v1/reports \\
        -H "Authorization: Bearer <token>" \\
        -H "ks-system-identification: 1.2.123456.2026" \\
        -H "Content-Type: application/json" \\
        -d @request.json | python3 -c "import sys,json; print(json.load(sys.stdin)['encryptedSessionKey'])" > aes_key.b64

    # 3. Pobranie zaszyfrowanego pliku z Azure Blob Storage
    az storage blob download \\
        --connection-string "<sas-connection-string>" \\
        --container-name rsr-output \\
        --name <ścieżka-do-pliku> \\
        --file report.enc

    # 4. Odszyfrowanie
    python3 decrypt_report.py \\
        --private-key rsr-data.key \\
        --aes-key aes_key.b64 \\
        --input report.enc \\
        --output report.parquet
"""

import argparse
import base64
import sys
from pathlib import Path

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


def load_private_key(path: str):
    pem = Path(path).read_bytes()
    return serialization.load_pem_private_key(pem, password=None)


def decrypt_aes_bundle(private_key, encrypted_b64: str) -> tuple[bytes, bytes]:
    encrypted = base64.b64decode(encrypted_b64.strip())
    bundle = private_key.decrypt(
        encrypted,
        padding.OAEP(
            mgf=padding.MGF1(algorithm=hashes.SHA256()),
            algorithm=hashes.SHA256(),
            label=None,
        ),
    )
    if len(bundle) != 44:
        raise ValueError(f"Nieprawidłowa długość klucza AES: {len(bundle)} B (oczekiwano 44 B)")
    return bundle[:32], bundle[32:]  # (aes_key, nonce)


def decrypt_file(aes_key: bytes, nonce: bytes, input_path: str, output_path: str):
    encrypted_data = Path(input_path).read_bytes()
    aesgcm = AESGCM(aes_key)
    plaintext = aesgcm.decrypt(nonce, encrypted_data, None)
    Path(output_path).write_bytes(plaintext)
    print(f"Odszyfrowano: {output_path} ({len(plaintext):,} bajtów)")


def main():
    parser = argparse.ArgumentParser(
        description="Deszyfruje raport retrospektywny zaszyfrowany algorytmem RSA/AESG/SHA256."
    )
    parser.add_argument("--private-key", required=True, metavar="PLIK",
                        help="Klucz prywatny RSA (PEM), para do publicKeyPem z żądania")
    parser.add_argument("--aes-key", required=True, metavar="PLIK_LUB_BASE64",
                        help="Zaszyfrowany klucz AES z odpowiedzi endpointu (plik .b64 lub base64 string)")
    parser.add_argument("--input", required=True, metavar="PLIK",
                        help="Zaszyfrowany plik raportu pobrany z Azure Blob Storage")
    parser.add_argument("--output", required=True, metavar="PLIK",
                        help="Docelowy plik po odszyfrowaniu (.parquet lub .csv)")
    args = parser.parse_args()

    aes_key_path = Path(args.aes_key)
    encrypted_b64 = aes_key_path.read_text().strip() if aes_key_path.exists() else args.aes_key

    private_key = load_private_key(args.private_key)
    aes_key, nonce = decrypt_aes_bundle(private_key, encrypted_b64)
    decrypt_file(aes_key, nonce, args.input, args.output)


if __name__ == "__main__":
    main()