add check on prime num

Author: marharytapaduchak

rsa.py

@@ -17,10 +17,42 @@ def egcd(a, b):
     g, y, x = egcd(b % a, a)
     return (g, x - (b // a) * y, y)
 
+def is_prime(n, k=5):
+    if n < 2:
+        return False
+    if n in (2, 3):
+        return True
+    if n % 2 == 0:
+        return False
+
+    # Записуємо n-1 як 2^s * d
+    s, d = 0, n - 1
+    while d % 2 == 0:
+        s += 1
+        d //= 2
+
+    for _ in range(k):
+        a = secrets.randbelow(n - 3) + 2  # [2, n-2]
+        x = pow(a, d, n)
+        if x in (1, n - 1):
+            continue
+        for _ in range(s - 1):
+            x = pow(x, 2, n)
+            if x == n - 1:
+                break
+        else:
+            return False
+    return True
+
+def generate_prime(bits=1024):
+    while True:
+        temp = secrets.randbits(bits)
+        if is_prime(temp):
+            return temp
+
 def rsa_algo():
     """RSA"""
-    # треба чекати чи просте число
-    num_1, num_2 = secrets.randbits(1024), secrets.randbits(1024)
+    num_1, num_2 = generate_prime(), generate_prime()
     n = num_1*num_2
 
     phi = (num_1-1)*(num_2-1)