{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[alice]: private: 895 -> public: 308\n",
      "[bob]: private: 804 -> public: 801\n",
      "[alice|bob] generated equal shared keys\n",
      "[alice] sent message 'Hello world' to bob\n",
      "[bob] received message 'Hello world' from alice\n"
     ]
    }
   ],
   "source": [
    "from fractions import Fraction\n",
    "import random\n",
    "from dataclasses import dataclass\n",
    "from sympy.ntheory.residue_ntheory import nthroot_mod\n",
    "from itertools import cycle\n",
    "\n",
    "@dataclass \n",
    "class ec:\n",
    "    a: int\n",
    "    b: int\n",
    "    p: int\n",
    "\n",
    "    def ec_add(self, p1, p2):\n",
    "        if all(v1 == v2 for v1, v2 in zip(p1, p2)):\n",
    "            #print(\"H\")\n",
    "            if self.p:\n",
    "                dy = Fraction((3*p1[0]**2 + self.a), 2*p1[1])\n",
    "            else:\n",
    "                dy =(3*p1[0]**2 + self.a) / (2*p1[1])\n",
    "        else:\n",
    "            if self.p:\n",
    "                dy = Fraction((p2[1] - p1[1]), (p2[0] - p1[0]))\n",
    "            else:\n",
    "                dy = (p2[1] - p1[1]) / (p2[0] - p1[0])\n",
    "        res_x = dy ** 2 - p1[0] - p2[0]\n",
    "        res_y = (dy * (p1[0] - res_x) - p1[1])\n",
    "        if self.p != None:\n",
    "            res_x = int(res_x % self.p) \n",
    "            res_y = int(res_y % self.p) \n",
    "            return (res_x, res_y)\n",
    "        return (float(res_x), float(res_y))\n",
    "    def ec_mul(self, p, s):\n",
    "        for _ in range(s):\n",
    "            p = self.ec_add(p, p)\n",
    "        return p\n",
    "    def __getitem__(self, pos):\n",
    "        x = pos[0]\n",
    "        x_given = type(x) != slice\n",
    "        if x_given: value = x**3+self.a*x+self.b\n",
    "        else: raise ValueError()\n",
    "        solutions = [(x % self.p, s) for s in nthroot_mod(value, 2, self.p, True)]\n",
    "        return solutions\n",
    "    \n",
    "@dataclass\n",
    "class Person:\n",
    "    name: str \n",
    "    private: tuple[int, int] \n",
    "    public: tuple[int, int] \n",
    "    shared: tuple[int, int] \n",
    "\n",
    "    def gen_keys(self, start, curve: ec):\n",
    "        self.private = random.randint(0, curve.p)\n",
    "        self.public = curve.ec_mul(start, self.private)\n",
    "        print(f\"[{self.name}]: private: {self.private} -> public: {self.public[1]}\")\n",
    "        return \n",
    "    \n",
    "    def gen_shared(self, public, curve: ec):\n",
    "        self.shared = curve.ec_mul(public, self.private)\n",
    "        return\n",
    "    \n",
    "    def diffie_hellman(curve: ec, alice, bob):\n",
    "        gen = (4, 10)\n",
    "        alice.gen_keys(gen, curve)\n",
    "        bob.gen_keys(gen, curve)\n",
    "        alice.gen_shared(bob.public, curve)\n",
    "        bob.gen_shared(alice.public, curve)\n",
    "        assert(alice.shared == bob.shared)\n",
    "        print(f\"[{alice.name}|{bob.name}] generated equal shared keys\")\n",
    "\n",
    "    def xor(data, key): \n",
    "        key = str(key)\n",
    "        data = data.encode() if isinstance(data, str) else data\n",
    "        key = key.encode() if isinstance(key, str) else key\n",
    "        return bytearray(a ^ b for a, b in zip(data, cycle(key)))\n",
    "\n",
    "    def send(self, msg, target):\n",
    "        encrypted = Person.xor(msg, self.shared[1])\n",
    "        print(f\"[{self.name}] sent message '{msg}' to {target.name}\")\n",
    "        target.recv(encrypted, self)\n",
    "\n",
    "    def recv(self, msg, source):\n",
    "        decrypted = str(Person.xor(msg, self.shared[1]), encoding='utf-8')\n",
    "        print(f\"[{self.name}] received message '{decrypted}' from {source.name}\")\n",
    "\n",
    "   \n",
    "if __name__ == \"__main__\":\n",
    "    curve = ec(0, 7, 1109)\n",
    "    alice = Person(\"alice\", 0, 0, 0)\n",
    "    bob = Person(\"bob\", 0, 0, 0)\n",
    "\n",
    "    Person.diffie_hellman(curve, alice, bob)\n",
    "    alice.send(\"Hello world\", bob)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "3.10.15",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}