src/generate.py

from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister
from qiskit_aer import Aer
from qiskit.visualization import plot_histogram
import numpy as np
import random
from typing import List, Tuple, Dict
import logging

class BB84Protocol:
    def __init__(self, key_length: int = 128):
        """
        Inicializa o protocolo BB84 com o tamanho desejado da chave.
        
        Args:
            key_length (int): Tamanho desejado da chave final em bits
        """
        self.key_length = key_length
        self.simulator = Aer.get_backend('qasm_simulator')
        
        # Configuração do logging
        logging.basicConfig(level=logging.INFO)
        self.logger = logging.getLogger(__name__)

    def _prepare_qubit(self, bit: int, basis: str) -> QuantumCircuit:
        """
        Prepara um único qubit no estado e base especificados.
        
        Args:
            bit (int): Bit a ser codificado (0 ou 1)
            basis (str): Base de medição ('Z' para rectilinear ou 'X' para diagonal)
            
        Returns:
            QuantumCircuit: Circuito quântico preparado
        """
        qr = QuantumRegister(1)
        cr = ClassicalRegister(1)
        qc = QuantumCircuit(qr, cr)
        
        # Prepara o estado do qubit
        if bit == 1:
            qc.x(qr[0])
            
        # Aplica a transformação de base se necessário
        if basis == 'X':
            qc.h(qr[0])
            
        return qc

    def _measure_qubit(self, circuit: QuantumCircuit, basis: str) -> QuantumCircuit:
        """
        Adiciona a medição em uma base específica ao circuito.
        
        Args:
            circuit (QuantumCircuit): Circuito a ser medido
            basis (str): Base de medição ('Z' ou 'X')
            
        Returns:
            QuantumCircuit: Circuito com medição adicionada
        """
        if basis == 'X':
            circuit.h(0)
        circuit.measure([0], [0])
        return circuit

    def generate_random_bits(self, n: int) -> List[int]:
        """
        Gera uma lista de bits aleatórios.
        
        Args:
            n (int): Número de bits a serem gerados
            
        Returns:
            List[int]: Lista de bits aleatórios
        """
        return [random.randint(0, 1) for _ in range(n)]

    def generate_random_bases(self, n: int) -> List[str]:
        """
        Gera uma lista de bases aleatórias.
        
        Args:
            n (int): Número de bases a serem geradas
            
        Returns:
            List[str]: Lista de bases aleatórias ('Z' ou 'X')
        """
        return [random.choice(['Z', 'X']) for _ in range(n)]

    def simulate_transmission(self, bits: List[int], bases: List[str]) -> Tuple[List[int], List[str]]:
        """
        Simula a transmissão dos qubits de Alice para Bob.
        
        Args:
            bits (List[int]): Bits a serem transmitidos
            bases (List[str]): Bases usadas para codificação
            
        Returns:
            Tuple[List[int], List[str]]: Medições de Bob e bases escolhidas
        """
        bob_bases = self.generate_random_bases(len(bits))
        bob_measurements = []
        
        for i in range(len(bits)):
            # Alice prepara o qubit
            circuit = self._prepare_qubit(bits[i], bases[i])
            
            # Bob mede o qubit
            circuit = self._measure_qubit(circuit, bob_bases[i])
            
            # Executa a simulação usando a nova API do Qiskit
            job = self.simulator.run(circuit, shots=1)
            result = job.result()
            counts = result.get_counts()
            measured_bit = int(list(counts.keys())[0])
            bob_measurements.append(measured_bit)
            
        return bob_measurements, bob_bases

    def sift_key(self, alice_bits: List[int], alice_bases: List[str], 
                 bob_measurements: List[int], bob_bases: List[str]) -> Tuple[List[int], List[int]]:
        """
        Realiza o processo de peneiramento da chave.
        
        Args:
            alice_bits (List[int]): Bits originais de Alice
            alice_bases (List[str]): Bases usadas por Alice
            bob_measurements (List[int]): Medições de Bob
            bob_bases (List[str]): Bases usadas por Bob
            
        Returns:
            Tuple[List[int], List[int]]: Chaves peneiradas de Alice e Bob
        """
        alice_key = []
        bob_key = []
        
        for i in range(len(alice_bits)):
            if alice_bases[i] == bob_bases[i]:
                alice_key.append(alice_bits[i])
                bob_key.append(bob_measurements[i])
                
        return alice_key, bob_key

    def check_eavesdropping(self, alice_key: List[int], bob_key: List[int], 
                           sample_size: float = 0.25) -> bool:
        """
        Verifica a presença de um espião comparando uma amostra das chaves.
        
        Args:
            alice_key (List[int]): Chave de Alice
            bob_key (List[int]): Chave de Bob
            sample_size (float): Proporção da chave a ser verificada
            
        Returns:
            bool: True se não houver evidência de espionagem, False caso contrário
        """
        if len(alice_key) != len(bob_key):
            return False
            
        sample_length = int(len(alice_key) * sample_size)
        indices = random.sample(range(len(alice_key)), sample_length)
        
        differences = 0
        for i in indices:
            if alice_key[i] != bob_key[i]:
                differences += 1
                
        error_rate = differences / sample_length
        self.logger.info(f"Taxa de erro na amostra: {error_rate:.2%}")
        
        # Define um limite aceitável de erro (por exemplo, 5%)
        return error_rate < 0.05

    def generate_key(self) -> Dict[str, List[int]]:
        """
        Executa o protocolo BB84 completo para gerar uma chave compartilhada.
        
        Returns:
            Dict[str, List[int]]: Dicionário contendo as chaves finais de Alice e Bob
        """
        # Fase 1: Preparação
        raw_bits = self.generate_random_bits(self.key_length * 4)  # Geramos mais bits para compensar o peneiramento
        alice_bases = self.generate_random_bases(len(raw_bits))
        
        self.logger.info(f"Iniciando protocolo BB84 para gerar chave de {self.key_length} bits")
        
        # Fase 2: Transmissão
        bob_measurements, bob_bases = self.simulate_transmission(raw_bits, alice_bases)
        
        # Fase 3: Peneiramento
        alice_key, bob_key = self.sift_key(raw_bits, alice_bases, bob_measurements, bob_bases)
        
        self.logger.info(f"Chave peneirada gerada com {len(alice_key)} bits")
        
        # Fase 4: Detecção de espionagem
        if not self.check_eavesdropping(alice_key, bob_key):
            self.logger.error("Possível tentativa de espionagem detectada!")
            return {"alice": [], "bob": []}
        
        # Fase 5: Finalização
        final_key_length = min(self.key_length, len(alice_key))
        alice_final_key = alice_key[:final_key_length]
        bob_final_key = bob_key[:final_key_length]
        
        self.logger.info(f"Protocolo BB84 concluído com sucesso. Chave final de {final_key_length} bits gerada.")
        
        return {
            "alice": alice_final_key,
            "bob": bob_final_key
        }

# Exemplo de uso
if __name__ == "__main__":
    # Criar uma instância do protocolo
    bb84 = BB84Protocol(key_length=64)
    
    # Gerar uma chave compartilhada
    result = bb84.generate_key()
    
    # Verificar se as chaves são idênticas
    if result["alice"] and result["bob"]:
        print("\nChave de Alice:", "".join(map(str, result["alice"])))
        print("Chave de Bob:  ", "".join(map(str, result["bob"])))
        print("\nAs chaves são idênticas?", result["alice"] == result["bob"])
        print("Tamanho da chave:", len(result["alice"]), "bits")
    else:
        print("Falha na geração da chave - possível tentativa de espionagem detectada.")