bb84/bb84.py at main · peterkim01/bb84 · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108

import numpy as np
import matplotlib.pyplot as plt
from qiskit import QuantumCircuit, transpile
from qiskit_aer import AerSimulator

def simulate_eve_measurement(alice_bit, alice_basis, eve_basis):
    # Eve gets correct result if her basis matches Alice's
    if alice_basis == eve_basis:
        return alice_bit
    else:
        return np.random.randint(2)

def simulate(eavesdropping):
    n = 100
    alice_bits = np.random.randint(2, size=n)
    alice_bases = np.random.randint(2, size=n)
    bob_bases = np.random.randint(2, size=n)
    eve_bases = np.random.randint(2, size=n) if eavesdropping else None

    circuits = []

    for i in range(n):
        qc = QuantumCircuit(1, 1)

        # alice encodes
        if alice_bits[i] == 1:
            qc.x(0)
        if alice_bases[i] == 1:
            qc.h(0)

        # eve intercepts and resends
        if eavesdropping:
            # Eve measures in her basis
            if eve_bases[i] == 1:
                qc.h(0)
            qc.measure(0, 0)

            eve_result = simulate_eve_measurement(alice_bits[i], alice_bases[i], eve_bases[i])
            eve_basis = eve_bases[i]

            # Re-encode qubit according to Eve's result and basis
            qc = QuantumCircuit(1, 1)
            if eve_result == 1:
                qc.x(0)
            if eve_basis == 1:
                qc.h(0)

        # bob applies his basis before measurement
        if bob_bases[i] == 1:
            qc.h(0)
        qc.measure(0, 0)

        circuits.append(qc)

    #SIMULATION
    backend = AerSimulator()
    transpiled_circuits = transpile(circuits, backend=backend)
    job = backend.run(transpiled_circuits, shots=1)
    results = job.result()
    bob_results = [int(list(results.get_counts(c).keys())[0]) for c in circuits]

    sifted_indices = np.where(alice_bases == bob_bases)[0]
    alice_key = alice_bits[sifted_indices]
    bob_key = np.array(bob_results)[sifted_indices]

    matches = np.sum(alice_key == bob_key)
    agreement_rate = matches / len(alice_key) if len(alice_key) > 0 else 0
    print(f"Key agreement rate: {agreement_rate * 100:.2f}%")

    #======= Visualization ========

    fig, axs = plt.subplots(3, 1, figsize=(12, 10))

    # Plot 1: Bases comparison
    axs[0].plot(alice_bases, label='Alice Bases', drawstyle='steps-mid', alpha=0.7)
    axs[0].plot(bob_bases, label='Bob Bases', drawstyle='steps-mid', alpha=0.7)
    if eavesdropping:
        axs[0].plot(eve_bases, label='Eve Bases', drawstyle='steps-mid', alpha=0.5)
    axs[0].set_title("Bases Used")
    axs[0].set_ylim(-0.5, 1.5)
    axs[0].legend()

    # Plot 2: Sifted Key Comparison
    axs[1].plot(alice_key, label='Alice Key (Sifted)', drawstyle='steps-mid', alpha=0.7)
    axs[1].plot(bob_key, label='Bob Key (Sifted)', drawstyle='steps-mid', alpha=0.7)
    axs[1].set_title("Sifted Key Bits: Alice vs Bob")
    axs[1].set_ylim(-0.5, 1.5)
    axs[1].legend()

    # Plot 3: Agreement Rate
    axs[2].bar(['Agreement', 'Disagreement'],
               [matches, len(alice_key) - matches],
               color=['green', 'red'])
    axs[2].set_title(f"Key Agreement Rate: {agreement_rate * 100:.2f}%")
    axs[2].set_ylabel("Number of Bits")

    plt.tight_layout()
    plt.show()

def main():
    eavesdropping = True if input("Enable eavesdropping? (T/F): ").strip().lower() == 't' else False

    print("Eavesdropping detected!" if eavesdropping else "No eavesdropping detected.")
    simulate(eavesdropping)

if __name__ == "__main__":
    main()