Bob.py

import socket
import base64

from cryptography.fernet import Fernet
from qkd import BB84

QHOST = '127.0.0.1'
QPORT = 12345

CHOST = '127.0.0.1'
CPORT = 12346

class ServerEntity():
    name = None
    counterparty = None
    counterparty_photon_stream = None
    counterparty_basis_stream = None
    qkd = None
    quantum_dict = None

    secret_key = None
    accepted_photon_stream = []
    accepted_basis_stream = []

    def __init__(self, name, counterparty):
        self.name = name
        self.counterparty = counterparty
        self.qkd = BB84()

    def handle_quantum_response(self, conn):
        try:
            print(f"{self.name}(server)> Connected with {self.counterparty} on Quantum Channel!")
            while(True):
                data = conn.recv(1024)
                if not data:
                    break
                photon_stream = data.decode().strip()
                print(f'{self.name}(server)> Received photons from {self.counterparty}: {photon_stream}')
                self.counterparty_photon_stream = photon_stream

                if photon_stream == "!/-\\":
                    reply_photon_stream = "!/-\\"
                    print(f'{self.name}(server)> Ending Quantum connection with {self.counterparty}: {reply_photon_stream}')
                    conn.sendall(reply_photon_stream.encode())
                    break
                else:
                    measurement_dict = self.qkd.measure_photons(photon_stream)
                    reply_photon_stream = measurement_dict["photon_stream"]
                    print(f'{self.name}(server)> Sending measured photons to {self.counterparty}: {reply_photon_stream}')   
                    conn.sendall(reply_photon_stream.encode())     
                    reply_photon_stream = "!/-\\"
                    print(f'{self.name}(server)> Ending Quantum connection with {self.counterparty}: {reply_photon_stream}')
                    conn.sendall(reply_photon_stream.encode())

            conn.close()
        except Exception as ex:
            print(f'{self.name}(server)> Quantum channel closed on exception:{Exception} {ex}.' )

        self.quantum_dict = measurement_dict

    def quantum_exchange_secret_key(self):
        try:
            # conn = True
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                    s.bind((QHOST, QPORT))
                    s.listen(1)
                    print(f"{self.name}(server)> initialising Quantum channel...")
                    while True:
                        conn, addr = s.accept()
                        self.handle_quantum_response(conn)

                        if self.counterparty_photon_stream:
                            break
                        
        except Exception as ex:
            print(f'{self.name}(server)> Exception on Quantum channel. {Exception}:{ex}')


    def handle_classical_response(self, conn):
        try:
            print(f"{self.name}(server)> Connected with {self.counterparty} on classical Channel!")
            while(True):
                data = conn.recv(1024)
                if not data:
                    break
                basis_stream = data.decode().strip()
                print(f'{self.name}(server)> Received basis stream from {self.counterparty}: {basis_stream}')
                self.counterparty_basis_stream = basis_stream

                if basis_stream == "END":
                    reply_stream = "END"
                    print(f'{self.name}(server)> Ending classical connection with {self.counterparty}: {reply_stream}')
                    conn.sendall(reply_stream.encode())
                    break
                else:
                    reply_stream =''.join(self.quantum_dict["basis_stream"])
                    print(f'{self.name}(server)> Sending basis to {self.counterparty}: {reply_stream}')   
                    conn.sendall(reply_stream.encode())     
                    reply_stream = "END"
                    print(f'{self.name}(server)> Ending classical connection with {self.counterparty}: {reply_stream}')
                    conn.sendall(reply_stream.encode())

            conn.close()
        except Exception as ex:
            print(f'{self.name}(server)> Quantum channel closed on exception:{Exception} {ex}.' )


    def classical_exchange_basis(self):
        try:
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                    s.bind((CHOST, CPORT))
                    s.listen(1)
                    print(f"{self.name}(server)> initialising classical channel...")
                    while True:
                        conn, _ = s.accept()
                        self.handle_classical_response(conn)

                        if self.counterparty_basis_stream:
                            break

                    print(f'{self.name}(server)> Establishing secret key.')
                    self.establish_secret_key()

                    print(f'{self.name}(server)> ESTABLISHED SECRET KEY: {self.secret_key}.')
        except Exception as ex:
            print(f'{self.name}(server)> Exception on classical channel. {Exception}:{ex}')

    def decrypt_data(self):
        _secret_key = self.secret_key[:32]
        ascii_key = _secret_key.encode('ascii')
        b64_ascii_key = base64.urlsafe_b64encode(ascii_key)
        f = Fernet(b64_ascii_key)
        self.data = f.decrypt(self.encrypted_data).decode('utf-8')
    
    def handle_encrypted_data(self, conn):
        try:
            print(f"{self.name}(server)> Connected with {self.counterparty} on classical channel for encrypted communications!")
            while(True):
                data = conn.recv(1024)
                if not data:
                    break
                
                encrypted_data = data.decode().strip()
                print(f'{self.name}(server)> Received encrypted data from {self.counterparty}: {encrypted_data}')

                self.encrypted_data = encrypted_data
                self.decrypt_data()
                
                if encrypted_data == "END":
                    reply_stream = "END"
                    print(f'{self.name}(server)> Ending classical connection with {self.counterparty}: {reply_stream}')
                    conn.sendall(reply_stream.encode())
                    break
                else:                     
                    reply_stream = "END"
                    print(f'{self.name}(server)> Ending classical connection with {self.counterparty}: {reply_stream}')
                    conn.sendall(reply_stream.encode())

            conn.close()
        except Exception as ex:
            print(f'{self.name}(server)> Classical channel closed on exception:{Exception} {ex}.' )


    def receive_encrypted_data(self):
        try:
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                    s.bind((CHOST, CPORT))
                    s.listen(1)
                    print(f"{self.name}(server)> initialising classical channel for encrypted transmission...")
                    while True:
                        conn, _ = s.accept()
                        self.handle_encrypted_data(conn)

                        if self.counterparty_basis_stream:
                            break
                    print(f'{self.name}(server)> Decrypted data with key calculated from QKD: {self.data}')
                
        except Exception as ex:
            print(f'{self.name}(server)> Exception on classical channel. {Exception}:{ex}')


    def establish_secret_key(self):
        basis_stream = ''.join(self.quantum_dict["basis_stream"])
        ix = 0
        for c in self.counterparty_photon_stream:
            current_photon = self.quantum_dict["photon_stream"][ix]
            if current_photon == c:
                ix_counterparty_basis = self.counterparty_basis_stream[ix]
                if ix_counterparty_basis ==  basis_stream[ix]:
                    self.accepted_photon_stream.append(c)
                    self.accepted_basis_stream.append(ix_counterparty_basis)
            ix += 1

        self.secret_key = self.qkd.decode_photons(self.accepted_photon_stream, self.accepted_basis_stream)


if __name__ == "__main__":
    bob = ServerEntity("Bob", "Alice")
    bob.quantum_exchange_secret_key()
    bob.classical_exchange_basis()
    bob.receive_encrypted_data()