Complete Demo

server.py

import sqlite3
import random
import os
import base64
import uuid
from typing import Tuple
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import io
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend

app = FastAPI(title="Simple Polynomial Key Server (String & CSV)")

class PolynomialMultiKeyEncryption:
    def __init__(self, raw_data: bytes = None):
        self.prime = 2**127 - 1
        if raw_data:
            self.model_bytes = raw_data
            self.secret_value = random.randint(1000, 9999)
            self.coefficients = [self.secret_value] + [random.randint(1, self.prime-1) for _ in range(2)]
            self.salt = os.urandom(16)
            self.encryption_key = self._derive_encryption_key(self.secret_value)
            self.encrypted_data = self._encrypt_data()
            
    def _derive_encryption_key(self, secret: int) -> bytes:
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=self.salt,
            iterations=100000
        )
        return kdf.derive(secret.to_bytes(16, 'big'))
    
    def _encrypt_data(self) -> bytes:
        iv = os.urandom(16)
        cipher = Cipher(algorithms.AES(self.encryption_key), modes.GCM(iv), backend=default_backend())
        encryptor = cipher.encryptor()
        encrypted = encryptor.update(self.model_bytes) + encryptor.finalize()
        return base64.b64encode(iv + encryptor.tag + encrypted)

    def decrypt_data(self) -> bytes:
        data = base64.b64decode(self.encrypted_data)
        iv, tag, encrypted = data[:16], data[16:32], data[32:]
        cipher = Cipher(algorithms.AES(self.encryption_key), modes.GCM(iv, tag), backend=default_backend())
        decryptor = cipher.decryptor()
        return decryptor.update(encrypted) + decryptor.finalize()

    def _evaluate_polynomial(self, x: int) -> int:
        x_mod = x % self.prime
        result = 0
        for power, coeff in enumerate(self.coefficients):
            term = coeff
            for _ in range(power):
                term = (term * x_mod) % self.prime
            result = (result + term) % self.prime
        return result

    def generate_new_key(self) -> str:
        x = random.randint(1, 1000000)
        y = self._evaluate_polynomial(x)
        return f"{x}-{y}"


class KeyManager:
    def __init__(self, db_name="demo_polynomial.db"):
        self.db_name = db_name
        self._create_tables()
        self.crypto_string = None
        self.crypto_csv = None
        self._initialize_demo_data()

    def _get_connection(self):
        return sqlite3.connect(self.db_name)

    def _create_tables(self):
        with self._get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("""
            CREATE TABLE IF NOT EXISTS UserIds (
                user_id INTEGER PRIMARY KEY AUTOINCREMENT,
                api_key TEXT UNIQUE,
                assigned_poly_key TEXT
            );
            """)
            cursor.execute("""
            CREATE TABLE IF NOT EXISTS decryptKeys (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                key_value TEXT UNIQUE
            );
            """)
            conn.commit()

    def _initialize_demo_data(self):
        string_demo = "Mensagem secreta mestre transmitida com seguranca!"
        self.crypto_string = PolynomialMultiKeyEncryption(string_demo.encode('utf-8'))
        print(self.crypto_string.encrypted_data.decode())
        csv_demo = "id,nome,cargo\n1,Alice,Engenheira\n2,Bruno,Cientista de Dados"
        self.crypto_csv = PolynomialMultiKeyEncryption(csv_demo.encode('utf-8'))

    def generate_and_register_keys(self, amount: int):
        new_keys = []
        for _ in range(amount):
            poly_key = self.crypto_string.generate_new_key()
            new_keys.append((poly_key,))
            
        with self._get_connection() as conn:
            cursor = conn.cursor()
            try:
                cursor.executemany("INSERT INTO decryptKeys (key_value) VALUES (?);", new_keys)
                conn.commit()
                return True
            except sqlite3.Error:
                conn.rollback()
                return False

    def get_or_assign_key(self, api_key: str):
        with self._get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT assigned_poly_key FROM UserIds WHERE api_key = ?;", (api_key,))
            existing = cursor.fetchone()
            if existing and existing[0]:
                return existing[0]

            cursor.execute("SELECT id, key_value FROM decryptKeys ORDER BY id ASC LIMIT 1;")
            result = cursor.fetchone()
            if not result:
                return None
            
            key_id, poly_key_value = result
            cursor.execute("INSERT OR REPLACE INTO UserIds (api_key, assigned_poly_key) VALUES (?, ?);", (api_key, poly_key_value))
            cursor.execute("DELETE FROM decryptKeys WHERE id = ?;", (key_id,))
            conn.commit()
            return poly_key_value

    def check_key_valid(self, api_key: str, key_to_check: str) -> bool:
        with self._get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT assigned_poly_key FROM UserIds WHERE api_key = ?;", (api_key,))
            result = cursor.fetchone()
            return result and result[0] == key_to_check


manager = KeyManager()

class KeyRequest(BaseModel):
    apiKey: str
    db: str

class CheckRequest(BaseModel):
    apiKey: str 
    key_to_check: str
    db: str

class DecryptPayloadRequest(BaseModel):
    apiKey: str 
    decryptChave: str
    target_type: str  # 'string' ou 'csv'


@app.post("/admin/generate-stock")
def generate_stock(amount: int = 500):
    if manager.generate_and_register_keys(amount):
        return {"status": "success", "message": f"{amount} chaves polinomiais registradas."}
    raise HTTPException(status_code=500, detail="Erro ao gerar estoque.")

@app.post("/getDecryptKey")
def get_decrypt_key(req: KeyRequest):
    """Retur the next valid decryptKey"""
    poly_key = manager.get_or_assign_key(req.apiKey)
    if not poly_key:
        raise HTTPException(status_code=404, detail="Nenhuma chave disponivel no estoque.")
    return {"decryptChave": poly_key}

@app.post("/checkDecryptKeyValid")
def check_decrypt_key_valid(req: CheckRequest):
    """Valid the x-y against the user apiKey."""
    is_valid = manager.check_key_valid(req.apiKey, req.key_to_check)
    return {"valid": is_valid}

@app.post("/decryptElement")
def decrypt_element(req: DecryptPayloadRequest):
    ""decrypt test from a dummy csv file and string"""
    # 1. Validação de segurança básica da amarração apiKey -> chave
    if not manager.check_key_valid(req.apiKey, req.decryptChave):
        raise HTTPException(status_code=403, detail="Chave nao autorizada para esta API Key.")

    try:
        # 2. Validação matemática do polinômio x-y
        parts = req.decryptChave.split('-')
        x, y = int(parts[0]), int(parts[1])
        
        expected_y = manager.crypto_string._evaluate_polynomial(x)
        if y != expected_y:
            raise HTTPException(status_code=400, detail="Chave matematicamente invalida.")

        # 3. Entrega o dado solicitado decifrado na memória do servidor
        if req.target_type == "csv":
            decrypted_bytes = manager.crypto_csv.decrypt_data()
        else:
            decrypted_bytes = manager.crypto_string.decrypt_data()

        return {
            "status": "success",
            "type": req.target_type,
            "data": decrypted_bytes.decode('utf-8')
        }
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Erro na descriptografia: {str(e)}")