main.py

import os
import random
import argparse

from qiskit import QuantumCircuit, transpile
from qiskit_aer import AerSimulator

from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives.ciphers.aead import AESGCM

# Meassure all qubits in the circuit and return the measurement results as a list of bits
def get_measurement_result(circuit) -> list[str]:
    simulator = AerSimulator()
    circ = transpile(circuit, simulator)
    result = simulator.run(circ, shots=1, memory=True).result()
    memory = result.get_memory(circ)

    # Reverse the memory to match the order of qubits
    # So the order of measured bits is from qubit 0 to qubit n-1.
    # Measurement of qubit 0 is in the first element)
    measurement_results = list(reversed(memory[0]))
    return measurement_results

# Generate a list of random bits using quantum randomness
def quantum_random_number_generator(num_bits) -> list[str]:
    circ = QuantumCircuit(1, 1) # 1 qubit, 1 classical bit
    circ.h(0) # hadamard gate
    circ.measure(0, 0) # measure qubit

    simulator = AerSimulator()
    circ = transpile(circ, simulator)

    # Since we only have one qubit, we run the circuit num_bits times to get num_bits random bits
    result = simulator.run(circ, shots=num_bits, memory=True).result()
    memory = result.get_memory(circ)

    # Return the measurement results as a list of bits
    return memory

# Assuming the inputs are horizontal polarized photons | 0 >
def encode_qubits(num_bits, random_bits, polarization_bases) -> QuantumCircuit:
    qc = QuantumCircuit(num_bits, num_bits)
    for i in range(num_bits):
        if polarization_bases[i] == '0':
            # Rectilinear basis
            if random_bits[i] == '1':
                qc.x(i)  # Apply X gate to make it | 1 >
            else:
                # Do nothing, because the polarization already horizontal
                pass
        else:
            # Diagonal basis
            if random_bits[i] == '0':
                qc.h(i)  # Apply H gate to make it | + >
            else:
                qc.x(i)  # Apply X gate to make it | 1 >
                qc.h(i)  # Then apply H gate to make it | - >

    qc.barrier() # Prevent compiler optimizations across the barrier
    return qc

def eavesdrop_qubits(qc, polarization_bases) -> QuantumCircuit:
    for i in range(qc.num_qubits):
        if polarization_bases[i] == '1':
            # Diagonal basis
            qc.h(i)  # Apply H gate to switch basis

    qc.measure(range(qc.num_qubits), range(qc.num_qubits))

    # Reverse the measurement results to match qubit order
    measurement_results = get_measurement_result(qc)

    # Reset the circuits (set all qubits to | 0 >)
    qc.reset(range(qc.num_qubits))

    for i in range(qc.num_qubits):
        if polarization_bases[i] == '0':
            # Rectilinear basis
            if measurement_results[i] == '1':
                qc.x(i)  # Apply X gate to make it | 1 >
            else:
                # Do nothing, because the polarization already horizontal
                pass
        else:
            # Diagonal basis
            if measurement_results[i] == '0':
                qc.h(i)  # Apply H gate to make it | + >
            else:
                qc.x(i)  # Apply X gate to make it | 1 >
                qc.h(i)  # Then apply H gate to make it | - >

    qc.barrier() # Prevent compiler optimizations across the barrier
    return qc

def measure_qubits(qc, polarization_bases) -> tuple[QuantumCircuit, list[str]]:
    for i in range(qc.num_qubits):
        if polarization_bases[i] == '1':
            # Diagonal basis
            qc.h(i)  # Apply H gate to switch basis

    qc.measure(range(qc.num_qubits), range(qc.num_qubits))

    # Reverse the measurement results to match qubit order
    measurement_results = get_measurement_result(qc)

    qc.barrier() # Prevent compiler optimizations across the barrier
    return (qc, measurement_results)

def get_correct_bases(alice_bases, bob_bases) -> list[bool]:
    correct_bases = []
    for i in range(len(alice_bases)):
        if alice_bases[i] == bob_bases[i]:
            correct_bases.append(True)
        else:
            correct_bases.append(False)
    return correct_bases

def get_shared_key(correct_bases, measurement_results) -> list[str]:
    shared_key = []
    for i in range(len(correct_bases)):
        if correct_bases[i]:
            shared_key.append(measurement_results[i])
    return shared_key

def derive_key(secret, salt, info) -> bytes:
    hkdf = HKDF(
        algorithm=hashes.SHA256(),
        length=32, # For AES-256
        salt=salt,
        info=info,
    )
    return hkdf.derive(secret)

def main() -> None:
    # Parse command line arguments
    parser = argparse.ArgumentParser(description="Quantum Key Distribution")
    parser.add_argument("num_bits", type=int, help="Number of bits to generate")
    parser.add_argument("--eavesdrop", help="Enable eavesdropping simulation (1 for yes, 0 for no)", type=int, choices=[0, 1], default=0)
    args = parser.parse_args()

    # alice is the initiator, bob is the receiver/responder
    num_bits = args.num_bits
    random_bits = quantum_random_number_generator(num_bits)

    # Generate random polarization base
    alice_bases = quantum_random_number_generator(num_bits)
    bob_bases = quantum_random_number_generator(num_bits)

    if args.eavesdrop:
        eve_bases = quantum_random_number_generator(num_bits)

    # Encode qubits based on alice's random bits and bases
    qc = encode_qubits(num_bits, random_bits, alice_bases)

    if args.eavesdrop:
        # If eavesdropping is enabled, Eve intercepts the qubits
        qc = eavesdrop_qubits(qc, eve_bases)

    # Bob measures the qubits based on his bases
    qc, measurement_results = measure_qubits(qc, bob_bases)
    print(qc.draw())

    # Alice and bob publish their bases over a public channel
    # Then they compare and keep only the bits where their bases matched
    correct_bases = get_correct_bases(alice_bases, bob_bases)

    # Generate shared key
    shared_key_alice = get_shared_key(correct_bases, random_bits)
    shared_key_bob = get_shared_key(correct_bases, measurement_results)

    # Detection of eavesdropping
    half_shared_key_bits_length = len(shared_key_alice) // 2
    alice_detection_bits = []
    bob_detection_bits = []

    # Alice generates random indices to check for eavesdropping
    # For this simulation, we give up half of the shared key bits for detection
    # So the key length for the encryption will be halved
    random.seed(os.urandom(16))
    for i in range(half_shared_key_bits_length):
        # Select a random index from the shared key bits
        new_shared_key_bits_length = len(shared_key_alice)
        idx = random.randint(0, new_shared_key_bits_length - 1)

        # Alice and Bob will remove the bit at the selected index from both shared keys
        alice_pop = shared_key_alice.pop(idx)
        bob_pop = shared_key_bob.pop(idx)
        alice_detection_bits.append(alice_pop)
        bob_detection_bits.append(bob_pop)

    # If the bits are altered, then we know that someone is eavesdropping
    if(alice_detection_bits != bob_detection_bits):
        print("Eavesdropper detected! Aborting key generation.")
        return

    # If no bits were altered, we can proceed to derive a shared secret key
    # Note that, there is still a small probability that Eve was eavesdropping but got lucky
    # that is, the detection bits were not altered.

    # We got a shared key, make it 32 bytes length for AES-256
    salt = os.urandom(16)
    info = b"shared_secret_key"
    hkdf_alice = derive_key(bytes(int(bit) for bit in shared_key_alice), salt, info)
    hkdf_bob = derive_key(bytes(int(bit) for bit in shared_key_bob), salt, info)
    print("Alice's key == Bob's key ? ", hkdf_alice == hkdf_bob)

    # Additional authenticated_data (AAD) and nonce can be publicly known,
    # typically, you send then along with the ciphertext.
    alice_encrypted_message = b"Secret Message"
    additional_authenticated_data = b'authenticated but unencrypted data'
    nonce = os.urandom(12)  # For AES GCM, a 12-byte nonce is recommended

    # Encrypt a message using the derived shared key.
    aesgcm_alice = AESGCM(hkdf_alice)
    ciphertext = aesgcm_alice.encrypt(nonce, alice_encrypted_message, additional_authenticated_data)
    print(f"Ciphertext (Hex): {ciphertext.hex()}")

    # There is two cases of undetected eavesdropping here (detection bits were not altered):
    # 1. Eve is eavesdropping, but got lucky that there is no bits that were altered on the actual shared key,
    # Thus the decryption will succeed.
    # 2. Eve is eavesdropping, but there is at least one bit altered on the actual shared key,
    # Thus the decryption will fail.

    aesgcm_bob = AESGCM(hkdf_bob)
    try:
        decrypted_message = aesgcm_bob.decrypt(nonce, ciphertext, additional_authenticated_data)
        print(f"Decrypted Message: {decrypted_message.decode()}")
    except:
        print(f"Decryption failed!")
        print("There is a potential eavesdropper, but alice and bob did not catch it!")
        print("Try to increase the number of detection bits!")
        return
    
    if args.eavesdrop:
        # This is the ideal scenario for Eve (we don't want this to happen)
        print("WARNING: There is an eavesdropper, but no bits were altered!")

if __name__ == "__main__":
    main()