From 22daf29b6f5d749f6216f7a83936993b9862d322 Mon Sep 17 00:00:00 2001
From: Alexander Munch-Hansen <alexmunchhansen@gmail.com>
Date: Wed, 2 Oct 2019 20:07:34 +0200
Subject: [PATCH] lmao

---
 week4.py | 23 ++++++++++++++++++++---
 week6.py |  6 +++---
 2 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/week4.py b/week4.py
index d60c899..8a31888 100644
--- a/week4.py
+++ b/week4.py
@@ -129,7 +129,7 @@ def is_prime(n: int, k: int) -> bool:
     return True
 
 
-def gen_prime(b: int, k: int = 10) -> int:
+def gen_prime(b: int, k: int = 4) -> int:
     """
     Generate strong probable prime by drawing integers at random until one passes the is_prime test.
     Adapted from pseudo-code at https://en.wikipedia.org/wiki/Miller%E2%80%93Rabin_primality_test.
@@ -142,14 +142,31 @@ def gen_prime(b: int, k: int = 10) -> int:
         n = random.randint(2**(b-1), (2**b)-1)
         if n % 2 == 0:
             continue
-        if is_prime(n, k):
+        # Check that the future value of q is prime
+        if is_prime(n, k) and is_prime(2*n+1, k):
             return n
 
 
+def find_primitive_root(p):
+    if p == 2:
+        return 1
+
+    p1 = 2
+    p2 = (p - 1) // p1
+
+    # test random g's until one is found that is a primitive root mod p
+    while True:
+        g = SystemRandom().randint(2, p - 1)
+
+        if not (pow(g, (p - 1) // p1, p) == 1):
+            if not pow(g, (p - 1) // p2, p) == 1:
+                return g
+
+
 def run(donor: BloodType, recipient: BloodType):
     p = gen_prime(128)
     q = 2 * p + 1
-    g = SystemRandom().randint(2, q)
+    g = find_primitive_root(p)
     #print("p:", p, "q:", q, "g:", g)
 
     elgamal = ElGamal(g, q, p)
diff --git a/week6.py b/week6.py
index b3231d2..d2f0c9f 100644
--- a/week6.py
+++ b/week6.py
@@ -8,7 +8,7 @@ from secrets import SystemRandom
 from typing import List, Union
 
 from .week1 import BloodType, blood_cell_compatibility_lookup
-from .week4 import gen_prime
+from .week4 import gen_prime, find_primitive_root
 
 
 class ElGamal:
@@ -220,8 +220,8 @@ class Bob:
 def run(da, db, ds, ra, rb, rs):
 
     p = gen_prime(256)
-    q = 2 * p + 1
-    g = SystemRandom().randint(2, q)
+    q = 2*p+1
+    g = find_primitive_root(p)
 
     elgamal = ElGamal(g, q, p)
     alice = Alice(ra=ra, rb=rb, rs=rs, elgamal=elgamal)