from fractions import Fraction
import random
from dataclasses import dataclass
from sympy.ntheory.residue_ntheory import nthroot_mod
from itertools import cycle

@dataclass 
class ec:
    a: int
    b: int
    p: int

    def ec_add(self, p1, p2):
        if all(v1 == v2 for v1, v2 in zip(p1, p2)):
            #print("H")
            if self.p:
                dy = Fraction((3*p1[0]**2 + self.a), 2*p1[1])
            else:
                dy =(3*p1[0]**2 + self.a) / (2*p1[1])
        else:
            if self.p:
                dy = Fraction((p2[1] - p1[1]), (p2[0] - p1[0]))
            else:
                dy = (p2[1] - p1[1]) / (p2[0] - p1[0])
        res_x = dy ** 2 - p1[0] - p2[0]
        res_y = (dy * (p1[0] - res_x) - p1[1])
        if self.p != None:
            res_x = int(res_x % self.p) 
            res_y = int(res_y % self.p) 
            return (res_x, res_y)
        return (float(res_x), float(res_y))
    def ec_mul(self, p, s):
        for _ in range(s):
            p = self.ec_add(p, p)
        return p
    def __getitem__(self, pos):
        x = pos[0]
        x_given = type(x) != slice
        if x_given: value = x**3+self.a*x+self.b
        else: raise ValueError()
        solutions = [(x % self.p, s) for s in nthroot_mod(value, 2, self.p, True)]
        return solutions
    
@dataclass
class Person:
    name: str 
    private: tuple[int, int] 
    public: tuple[int, int] 
    shared: tuple[int, int] 

    def gen_keys(self, start, curve: ec):
        self.private = random.randint(0, curve.p)
        self.public = curve.ec_mul(start, self.private)
        print(f"[{self.name}]: private: {self.private} -> public: {self.public[1]}")
        return 
    
    def gen_shared(self, public, curve: ec):
        self.shared = curve.ec_mul(public, self.private)
        return
    
    def diffie_hellman(curve: ec, alice, bob):
        gen = (4, 10)
        alice.gen_keys(gen, curve)
        bob.gen_keys(gen, curve)
        alice.gen_shared(bob.public, curve)
        bob.gen_shared(alice.public, curve)
        assert(alice.shared == bob.shared)
        print(f"[{alice.name}|{bob.name}] generated equal shared keys")

    def xor(data, key): 
        key = str(key)
        data = data.encode() if isinstance(data, str) else data
        key = key.encode() if isinstance(key, str) else key
        return bytearray(a ^ b for a, b in zip(data, cycle(key)))

    def send(self, msg, target):
        encrypted = Person.xor(msg, self.shared[1])
        print(f"[{self.name}] sent message '{msg}' to {target.name}")
        target.recv(encrypted, self)

    def recv(self, msg, source):
        decrypted = str(Person.xor(msg, self.shared[1]), encoding='utf-8')
        print(f"[{self.name}] received message '{decrypted}' from {source.name}")

   
if __name__ == "__main__":
    curve = ec(0, 7, 1109)
    alice = Person("alice", 0, 0, 0)
    bob = Person("bob", 0, 0, 0)

    Person.diffie_hellman(curve, alice, bob)
    alice.send("Hello world", bob)