101 lines
2.4 KiB
Python
101 lines
2.4 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import hashlib
|
||
|
|
import itertools
|
||
|
|
import random
|
||
|
|
import secrets
|
||
|
|
from typing import List
|
||
|
|
|
||
|
|
from week1 import BloodType
|
||
|
|
|
||
|
|
|
||
|
|
def sha256(b: bytes) -> bytes:
|
||
|
|
return hashlib.sha256(b).digest()
|
||
|
|
|
||
|
|
|
||
|
|
def rand_bytes():
|
||
|
|
return secrets.SystemRandom().getrandbits(128).to_bytes(16, "big")
|
||
|
|
|
||
|
|
|
||
|
|
def xor_bytes(a: bytes, b: bytes, k=256) -> bytes:
|
||
|
|
return (int.from_bytes(a, "big") ^ int.from_bytes(b, "big")).to_bytes(k, "big")
|
||
|
|
|
||
|
|
|
||
|
|
class Gate:
|
||
|
|
def __init__(self, left: Gate, right: Gate, index: int) -> None:
|
||
|
|
self.left = left
|
||
|
|
self.right = right
|
||
|
|
self.i = index
|
||
|
|
self.k = {
|
||
|
|
0: rand_bytes(),
|
||
|
|
1: rand_bytes()
|
||
|
|
}
|
||
|
|
self.output = None
|
||
|
|
|
||
|
|
c_prime = {}
|
||
|
|
for a, b in itertools.product((0, 1), repeat=2):
|
||
|
|
c_prime[(a, b)] = xor_bytes(
|
||
|
|
sha256(self.left.k[a] + self.right.k[b] + self.i.to_bytes(1, "big")),
|
||
|
|
self.k[not a * b] + bytes(128 // 8)
|
||
|
|
)
|
||
|
|
pi = list(itertools.product((0, 1), repeat=2))
|
||
|
|
random.shuffle(pi)
|
||
|
|
self.c = {i: c_prime[p]
|
||
|
|
for i, p in enumerate(pi)}
|
||
|
|
|
||
|
|
|
||
|
|
class InputWire:
|
||
|
|
def __init__(self, index) -> None:
|
||
|
|
self.i = index
|
||
|
|
self.k = {
|
||
|
|
0: rand_bytes(),
|
||
|
|
1: rand_bytes()
|
||
|
|
}
|
||
|
|
self.output = None
|
||
|
|
|
||
|
|
|
||
|
|
class Circuit:
|
||
|
|
def __init__(self, input_wires: List[InputWire], gates: List[Gate]) -> None:
|
||
|
|
self.input_wires = input_wires
|
||
|
|
self.gates = gates
|
||
|
|
|
||
|
|
def evaluate(self, x: List[bytes]) -> bytes:
|
||
|
|
for i, xi in enumerate(x):
|
||
|
|
self.input_wires[i].output = xi
|
||
|
|
|
||
|
|
for gate in self.gates:
|
||
|
|
for j in range(4):
|
||
|
|
xor = xor_bytes(
|
||
|
|
sha256(gate.left.output + gate.right.output + gate.i),
|
||
|
|
gate.c[j]
|
||
|
|
)
|
||
|
|
k_prime, tau = xor[:16], xor[16:]
|
||
|
|
if tau == bytes(16):
|
||
|
|
gate.output = k_prime
|
||
|
|
|
||
|
|
return self.gates[-1].output
|
||
|
|
|
||
|
|
|
||
|
|
def encode(e: List[InputWire], x: List[int]) -> List[bytes]:
|
||
|
|
return [e[i].k[xi]
|
||
|
|
for i, xi in enumerate(x)]
|
||
|
|
|
||
|
|
|
||
|
|
class Alice:
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
class Bob:
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
def run(donor: BloodType, recipient: BloodType):
|
||
|
|
m1 = Alice.choose(x)
|
||
|
|
m2 = Bob.transfer(y, m2)
|
||
|
|
z = Alice.retrieve(m2)
|
||
|
|
return z
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
run(donor=BloodType.A_NEGATIVE, recipient=BloodType.B_POSITIVE)
|