quantum/00_qkd.py

108 lines
2.8 KiB
Python

from dataclasses import dataclass, field
import random
STREAM_LENGTH = 100
@dataclass
class Photon(object):
basis: int
value: int
_measured: bool = False
_measured_value: bool = None
def measure(self, basis):
self._measured = True
if self.basis == basis:
return self.value
else:
if not self._measured:
self._measured_value = random.randint(0, 1)
return self._measured_value
def __repr__(self):
return '{}:{}'.format(self.basis, self.value)
@dataclass
class Alice:
_otp: list = None
bases: list = field(default_factory=list)
_values: list = field(default_factory=list)
def encode_qbit(self, value):
basis = random.randint(0, 1)
self.bases.append(basis)
return Photon(basis, value)
def generate_stream(self, stream_size=10):
for _ in range(stream_size):
qbit = random.randint(0, 1)
self._values.append(qbit)
yield self.encode_qbit(qbit)
def set_otp(self, other_correct):
self._otp = [v for c, v in zip(other_correct, self._values) if c]
def get_otp_length(self):
return len(self._otp)
@dataclass
class Bob:
correct: list = None
_otp: list = None
_bases: list = field(default_factory=list)
_values: list = field(default_factory=list)
def ingest_stream(self, stream):
for photon in stream:
basis = random.randint(0, 1)
self._bases.append(basis)
qbit_value = photon.measure(basis)
self._values.append(qbit_value)
def insecure_get_bases(self, other_bases):
self.correct = [not (a ^ b) for a, b in zip(self._bases, other_bases)]
self._otp = [v for c, v in zip(self.correct, self._values) if c]
@dataclass
class Eve:
_bases: list = field(default_factory=list)
_values: list = field(default_factory=list)
def ingest_stream(self, stream):
for photon in stream:
basis = random.randint(0, 1)
self._bases.append(basis)
qbit_value = photon.measure(basis)
self._values.append(qbit_value)
yield photon
def main():
alice, bob, eve = Alice(), Bob(), Eve()
# Step 1. Generate a stream of random photons polarizations in a random
# basis
stream = alice.generate_stream(STREAM_LENGTH)
# Step 1.5 EVE ALSO INGESTS AND COPIES THE PHOTONS
# stream = eve.ingest_stream(stream)
# Step 2. Bob ingests the stream of photons
bob.ingest_stream(stream)
bob.insecure_get_bases(alice.bases)
alice.set_otp(bob.correct)
assert alice._otp == bob._otp
print("OTP Length: {}".format(alice.get_otp_length()))
print("OTP Length: {}".format(alice._otp))
if __name__ == "__main__":
main()