import base64
import numpy as np
np.set_printoptions(linewidth=110)

from qiskit import execute
from qiskit_aer import AerSimulator

from bb84 import initialize_protocol, encode_qubits, measure_qubits, filter_qubits, array_to_string
from error_correction import LinearCode
from secret_utils import generate_token, convert_to_octets

Set simulation parameters

number_of_qubits = 50
backend = AerSimulator(method="stabilizer")

Protocol initialization

1. Create two random sequences of 0's and 1's for Alice:
   • one will determine the encoding basis for each qubit
   • the other one will determine the state of each qubit
2. create a random sequence for Bob
   • this will determine the measurement basis for each qubit

# For Alice, the important basis is encoding basis.
encoding_basis_A, states_A, _ = initialize_protocol(number_of_qubits)

# Print the initial state of Alice
sent_bits = array_to_string(states_A)
print(f'Alice will encode her bits: \n{states_A} \nusing the encoding bases: \n{encoding_basis_A}')

Alice will encode her bits: 
[1 0 1 0 1 0 1 0 1 0 0 0 1 0 0 1 1 1 0 0 1 0 1 1 0 1 1 1 1 1 1 0 1 0 1 1 1 0 0 1 1 0 1 1 0 1 1 0 0 1] 
using the encoding bases: 
[0 0 0 0 0 1 1 1 0 0 1 1 1 0 0 0 1 1 1 1 0 0 1 0 0 1 1 1 1 1 1 0 1 0 1 1 0 1 0 1 1 0 1 0 1 1 0 1 0 0]

# For Bob, the relevant basis is measurement basis.
_, _, measurement_basis_B = initialize_protocol(number_of_qubits)
print(f'Bob will measure using these bases: \n{measurement_basis_B}')

Bob will measure using these bases: 
[1 0 1 1 1 1 1 0 0 1 1 0 1 0 0 1 0 0 0 1 0 0 1 0 0 0 1 1 0 1 0 1 0 0 0 1 0 0 1 1 0 1 1 0 1 1 1 1 1 1]

Encoding the qubits

Alice encodes the values of her qubits using according bases from encoding_bases_A.
This is stored as a Qiskit quantum circuit.

encoded_qubits_A = encode_qubits(number_of_qubits, states_A, encoding_basis_A)
encoded_qubits_A.draw()

      ┌───┐     
 q_0: ┤ X ├─────
      └───┘     
 q_1: ──────────
      ┌───┐     
 q_2: ┤ X ├─────
      └───┘     
 q_3: ──────────
      ┌───┐     
 q_4: ┤ X ├─────
      ├───┤     
 q_5: ┤ H ├─────
      ├───┤┌───┐
 q_6: ┤ X ├┤ H ├
      ├───┤└───┘
 q_7: ┤ H ├─────
      ├───┤     
 q_8: ┤ X ├─────
      └───┘     
 q_9: ──────────
      ┌───┐     
q_10: ┤ H ├─────
      ├───┤     
q_11: ┤ H ├─────
      ├───┤┌───┐
q_12: ┤ X ├┤ H ├
      └───┘└───┘
q_13: ──────────

q_14: ──────────
      ┌───┐     
q_15: ┤ X ├─────
      ├───┤┌───┐
q_16: ┤ X ├┤ H ├
      ├───┤├───┤
q_17: ┤ X ├┤ H ├
      ├───┤└───┘
q_18: ┤ H ├─────
      ├───┤     
q_19: ┤ H ├─────
      ├───┤     
q_20: ┤ X ├─────
      └───┘     
q_21: ──────────
      ┌───┐┌───┐
q_22: ┤ X ├┤ H ├
      ├───┤└───┘
q_23: ┤ X ├─────
      └───┘     
q_24: ──────────
      ┌───┐┌───┐
q_25: ┤ X ├┤ H ├
      ├───┤├───┤
q_26: ┤ X ├┤ H ├
      ├───┤├───┤
q_27: ┤ X ├┤ H ├
      ├───┤├───┤
q_28: ┤ X ├┤ H ├
      ├───┤├───┤
q_29: ┤ X ├┤ H ├
      ├───┤├───┤
q_30: ┤ X ├┤ H ├
      └───┘└───┘
q_31: ──────────
      ┌───┐┌───┐
q_32: ┤ X ├┤ H ├
      └───┘└───┘
q_33: ──────────
      ┌───┐┌───┐
q_34: ┤ X ├┤ H ├
      ├───┤├───┤
q_35: ┤ X ├┤ H ├
      ├───┤└───┘
q_36: ┤ X ├─────
      ├───┤     
q_37: ┤ H ├─────
      └───┘     
q_38: ──────────
      ┌───┐┌───┐
q_39: ┤ X ├┤ H ├
      ├───┤├───┤
q_40: ┤ X ├┤ H ├
      └───┘└───┘
q_41: ──────────
      ┌───┐┌───┐
q_42: ┤ X ├┤ H ├
      ├───┤└───┘
q_43: ┤ X ├─────
      ├───┤     
q_44: ┤ H ├─────
      ├───┤┌───┐
q_45: ┤ X ├┤ H ├
      ├───┤└───┘
q_46: ┤ X ├─────
      ├───┤     
q_47: ┤ H ├─────
      └───┘     
q_48: ──────────
      ┌───┐     
q_49: ┤ X ├─────
      └───┘     

Next, Alice sends the qubits to Bob.
In reality, this is done via QKD link. In our simulator, Alice sends a serialized quantum circuit (*.qpy object).
Bob measures the qubits using his measurement basis.

measured_circuit = measure_qubits(encoded_qubits_A, measurement_basis_B)
result = execute(measured_circuit, backend=backend, shots=1).result()
measured_bits = list(result.get_counts(measured_circuit))[0][::-1]
measured_bits

'00110011110010010110101101111110100111010011011010'

Key sifting

After Bob has measured the qubits, he sends his measurement bases to Alice. She responds to Bob by sending him her encoding bases. Now both parties know both the encoding basis and measurement basis for each qubit.

In the key sifting phase, both sides keep only the qubits for which encoding basis and measurement basis were the same.

If there are no transmission or measurement errors, at this step both Alice and Bob should have the same key.

alice_raw_key = filter_qubits(sent_bits, encoding_basis_A, measurement_basis_B)
print(alice_raw_key)

[0 0 1 1 0 1 0 0 0 1 0 1 1 0 1 1 1 0 1 1 1 1 1 0 1 0]

bob_raw_key = filter_qubits(measured_bits, encoding_basis_A, measurement_basis_B)
print(bob_raw_key)

[0 0 1 1 0 1 0 0 0 1 0 1 1 0 1 1 1 0 1 1 1 1 1 0 1 0]

Now, we will introduce an error on purpose to demonstrate error correction.
We will flip the third bit in Bob's key to simulate a transmission error.

# Flip Bob's third bit for sake of error correction demonstration
bob_raw_key[2] = np.mod(bob_raw_key[2] + 1,2)
print(bob_raw_key)

[0 0 0 1 0 1 0 0 0 1 0 1 1 0 1 1 1 0 1 1 1 1 1 0 1 0]

Key reconciliation

After the key sifting phase, Alice and Bob proceed to key reconciliation.

To reconciliate the key, techniques of error correction are used. In our case - a linear error-correcting code.
The process is following:

1. Both sides create a linear code of required dimension (depending on the resulting raw key length).
2. Both sides calculate the generator matrix, parity check matrix and syndrome table.
3. Alice encodes her key into a codeword using the generator matrix.
4. Alice sends to Bob the parity check bits (all the bits that are not the key itself).
5. Bob concatenates the bits he received from Alice together with his key. This plays the role of a message with potential errors.
6. Bob uses the parity-check matrix and syndrome to decide if there are errors in his key (with respect to Alice's key) and to correct potential errors.

# First, create a linear code of required dimension (depending on the resulting raw key length).
code_dimension = len(alice_raw_key)
codeword_length = len(alice_raw_key) + 5
error_correcting_code = LinearCode(k=code_dimension, n=codeword_length)

generator_matrix = error_correcting_code.generator_matrix()
parity_check = error_correcting_code.get_parity_check_matrix()
syndrome_table = error_correcting_code.get_syndrome_decoding_table()

encoded_key_A = np.dot(alice_raw_key, generator_matrix)
encoded_key_A = np.mod(encoded_key_A, 2)
print(f'Key of Alice after encoding is: {encoded_key_A}')
syndrome_A = np.mod(np.dot(parity_check.T, encoded_key_A), 2)
print(f'Syndrom of Alice (should be all zero): {syndrome_A}')
to_send = encoded_key_A[code_dimension:]
print(f'Information sent to Bob: {to_send}')

Key of Alice after encoding is: [0 0 1 1 0 1 0 0 0 1 0 1 1 0 1 1 1 0 1 1 1 1 1 0 1 0 1 1 0 0 1]
Syndrom of Alice (should be all zero): [0 0 0 0 0]
Information sent to Bob: [1 1 0 0 1]

encoded_key_B = np.concatenate((bob_raw_key, to_send))
syndrome_B = np.mod(np.dot(parity_check.T, encoded_key_B), 2)
print(syndrome_B)

[1 1 0 1 1]

The syndrome on Bob's site is not all zeros, which means that there is an error as compared to Alice's key! To correct the error, we need to look up the received syndrome in the syndrome table and add the corresponding value to Bob's raw key.

correction_mask = syndrome_table[tuple(syndrome_B)]
corrected_key = np.mod(bob_raw_key + correction_mask[:code_dimension], 2)
print('Key of Alice:'.ljust(30) + f'{alice_raw_key}')
print(f'Key of Bob before correction: {bob_raw_key}')
print('Key of Bob after correction:'.ljust(30) + f'{corrected_key}')

Key of Alice:                 [0 0 1 1 0 1 0 0 0 1 0 1 1 0 1 1 1 0 1 1 1 1 1 0 1 0]
Key of Bob before correction: [0 0 0 1 0 1 0 0 0 1 0 1 1 0 1 1 1 0 1 1 1 1 1 0 1 0]
Key of Bob after correction:  [0 0 1 1 0 1 0 0 0 1 0 1 1 0 1 1 1 0 1 1 1 1 1 0 1 0]

Generate alphanumeric keys from the sequences of 0's and 1's

If you need keys that are made of letters (e.g. for tokens or other regular use), you should group the key into octets,
create a bytearray of octets and then encode each octet according to ASCII.

Please note that this shortens they key and to work properly quite a key of at least 8 bit is needed.

ASCII_key = base64.b64encode(convert_to_octets(array_to_string(corrected_key))).decode('ascii')
print(ASCII_key)

NFu+