Quantum_Cryptography/model/eve.py

from dataclasses import dataclass, field
import typing
from random import choice, uniform


@dataclass
class EveBS:
    hijack_probability: float = 0.45
    my_bases: typing.List[bool] = field(default_factory=list)
    obtained_data: typing.List[typing.Optional[bool]] = field(default_factory=list)

    def process_photons(self, photons: typing.List[typing.Tuple[bool, bool]]) -> typing.List[typing.Tuple[bool, bool]]:
        basis = choice([False, True])
        self.my_bases.append(basis)
        clicks = [0, 0]
        passed = []
        for photon in photons:
            if uniform(0, 1) <= self.hijack_probability:
                photon_bit, photon_basis = photon
                if photon_basis == basis:
                    clicks[photon_bit] += 1
                else:
                    clicks[choice([0, 1])] += 1
            else:
                passed.append(photon)
        if any(clicks) and not all(clicks):
            self.obtained_data.append(bool(clicks[1]))
        else:
            self.obtained_data.append(None)
        return passed


@dataclass
class EvePNS:
    k: int = 1
    my_bases: typing.List[bool] = field(default_factory=list)
    obtained_data: typing.List[typing.Optional[bool]] = field(default_factory=list)

    def process_photons(self, photons: typing.List[typing.Tuple[bool, bool]]) -> typing.List[typing.Tuple[bool, bool]]:
        basis = choice([False, True])
        self.my_bases.append(basis)
        clicks = [0, 0]
        passed = []
        captured = 0
        for photon in photons:
            if captured < self.k:
                captured += 1
                photon_bit, photon_basis = photon
                if photon_basis == basis:
                    clicks[photon_bit] += 1
                else:
                    clicks[choice([0, 1])] += 1
            else:
                passed.append(photon)
        if any(clicks) and not all(clicks):
            self.obtained_data.append(bool(clicks[1]))
        else:
            self.obtained_data.append(None)
        return passed