crypto_computing/week6.py

251 lines
6.7 KiB
Python
Raw Normal View History

2019-10-02 14:53:44 +00:00
from __future__ import annotations
import hashlib
import itertools
2019-10-02 17:38:54 +00:00
import math
2019-10-02 14:53:44 +00:00
import random
2019-10-02 17:38:54 +00:00
from secrets import SystemRandom
from typing import List, Union
from .week1 import BloodType, blood_cell_compatibility_lookup
2019-10-02 18:07:34 +00:00
from .week4 import gen_prime, find_primitive_root
2019-10-02 17:38:54 +00:00
class ElGamal:
def __init__(self, g, q, p):
self.gen_ = g
self.order = q
self.p = p
def gen_key(self):
key = SystemRandom().randint(1, self.order)
while math.gcd(self.order, key) != 1:
key = SystemRandom().randint(1, self.order)
return key
def gen(self, sk):
h = pow(self.gen_, sk, self.order)
return (self.gen_, h)
def enc(self, m, pk):
# sample random r \in Zq
r = SystemRandom().randint(1, self.order)
g, h = pk
s = pow(h, r, self.order)
p = pow(g, r, self.order)
tmp = int.from_bytes(m, "big")
c = s * tmp
return c, p
def dec(self, c, sk):
c1, c2 = c
h = pow(c2, sk, self.order)
m = c1 // h
return m.to_bytes(16, "big")
def ogen(self):
s = SystemRandom().randint(1, self.order)
h = pow(s, 2, self.order)
return self.gen_, h
2019-10-02 14:53:44 +00:00
def sha256(b: bytes) -> bytes:
return hashlib.sha256(b).digest()
def rand_bytes():
2019-10-02 17:38:54 +00:00
return SystemRandom().getrandbits(128).to_bytes(16, "big")
2019-10-02 14:53:44 +00:00
2019-10-02 17:38:54 +00:00
def xor_bytes(a: bytes, b: bytes, k=32) -> bytes:
2019-10-02 14:53:44 +00:00
return (int.from_bytes(a, "big") ^ int.from_bytes(b, "big")).to_bytes(k, "big")
class Gate:
2019-10-02 17:38:54 +00:00
def __init__(self, left: Union[Gate, InputWire], right: Union[Gate, InputWire], index: int) -> None:
2019-10-02 14:53:44 +00:00
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")),
2019-10-02 17:38:54 +00:00
self.k[self.f(a, b)] + bytes(16)
2019-10-02 14:53:44 +00:00
)
pi = list(itertools.product((0, 1), repeat=2))
random.shuffle(pi)
self.c = {i: c_prime[p]
for i, p in enumerate(pi)}
2019-10-02 17:38:54 +00:00
def f(self, a, b):
raise NotImplemented
class ImplyGate(Gate):
def f(self, a, b):
return a >= b
class AndGate(Gate):
def f(self, a, b):
return a * b
2019-10-02 14:53:44 +00:00
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
2019-10-02 17:38:54 +00:00
@property
def d(self):
return self.gates[-1].k
2019-10-02 14:53:44 +00:00
def evaluate(self, x: List[bytes]) -> bytes:
2019-10-02 17:38:54 +00:00
for i, input_wire in enumerate(self.input_wires):
input_wire.output = x[i]
2019-10-02 14:53:44 +00:00
for gate in self.gates:
for j in range(4):
xor = xor_bytes(
2019-10-02 17:38:54 +00:00
sha256(gate.left.output + gate.right.output + gate.i.to_bytes(1, "big")),
2019-10-02 14:53:44 +00:00
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:
2019-10-02 17:38:54 +00:00
def __init__(self, ra, rb, rs, elgamal):
self.elgamal = elgamal
self.sks = None
self.input = [ra, rb, rs]
self.keys = None
def send_pks(self):
pks = []
self.sks = []
for idx, input_ in enumerate(self.input):
sk = self.elgamal.gen_key()
pk = self.elgamal.gen(sk)
self.sks.append(sk)
fake_pk = self.elgamal.ogen()
pk_tuple = [fake_pk]
pk_tuple.insert(input_, pk)
pks.append(pk_tuple)
return pks
def retrieve(self, circuit, bob_keys, ciphers):
self.keys = []
for idx, sk in enumerate(self.sks):
self.keys.append(self.elgamal.dec(ciphers[idx][self.input[idx]], sk))
all_keys = self.keys + bob_keys
res = circuit.evaluate(all_keys)
if circuit.d[0] == res:
return 0
if circuit.d[1] == res:
return 1
raise Exception("Fuck you")
2019-10-02 14:53:44 +00:00
class Bob:
2019-10-02 17:38:54 +00:00
def __init__(self, da, db, ds, elgamal):
input_wire1 = InputWire(0)
input_wire2 = InputWire(1)
input_wire3 = InputWire(2)
input_wire4 = InputWire(3)
input_wire5 = InputWire(4)
input_wire6 = InputWire(5)
impl_gate_1 = ImplyGate(input_wire1, input_wire4, 6)
impl_gate_2 = ImplyGate(input_wire2, input_wire5, 7)
impl_gate_3 = ImplyGate(input_wire3, input_wire6, 8)
and_gate_1 = AndGate(impl_gate_1, impl_gate_2, 9)
and_gate_2 = AndGate(and_gate_1, impl_gate_3, 10)
self.circuit = Circuit(
input_wires=[input_wire1, input_wire2, input_wire3, input_wire4, input_wire5, input_wire6],
gates=[impl_gate_1, impl_gate_2, impl_gate_3, and_gate_1, and_gate_2]
)
self.own_keys = encode([input_wire4, input_wire5, input_wire6], [da, db, ds])
self.key_set = [x.k.values() for x in [input_wire1, input_wire2, input_wire3]]
self.elgamal = elgamal
self.pks = None
def receive_pks(self, pks):
self.pks = pks
def transfer_messages(self):
ciphers = []
for idx, (k0, k1) in enumerate(self.key_set):
pk0, pk1 = self.pks[idx]
c0 = self.elgamal.enc(k0, pk0)
c1 = self.elgamal.enc(k1, pk1)
ciphers.append((c0, c1))
return self.circuit, self.own_keys, ciphers
def run(da, db, ds, ra, rb, rs):
p = gen_prime(256)
2019-10-02 18:07:34 +00:00
q = 2*p+1
g = find_primitive_root(p)
2019-10-02 17:38:54 +00:00
elgamal = ElGamal(g, q, p)
alice = Alice(ra=ra, rb=rb, rs=rs, elgamal=elgamal)
bob = Bob(da=da, db=db, ds=ds, elgamal=elgamal)
2019-10-02 14:53:44 +00:00
2019-10-02 17:38:54 +00:00
bob.receive_pks(alice.send_pks())
pls = alice.retrieve(*bob.transfer_messages())
2019-10-02 14:53:44 +00:00
2019-10-02 17:38:54 +00:00
return pls
2019-10-02 14:53:44 +00:00
def main():
2019-10-02 17:38:54 +00:00
green = 0
red = 0
for i, recipient in enumerate(BloodType):
for j, donor in enumerate(BloodType):
z = run(*donor.value, *recipient.value)
lookup = blood_cell_compatibility_lookup(recipient, donor)
if lookup == z:
green += 1
else:
print(f"'{BloodType(donor).name} -> {BloodType(recipient).name}' should be {lookup}.")
red += 1
print("Green:", green)
print("Red :", red)
# run(donor=BloodType.A_NEGATIVE, recipient=BloodType.B_POSITIVE)