Improve the key generator

1 files changed, 35 insertions, 47 deletions

diff --git a/forge_key.py b/forge_key.py
index 3bed393..756663b 100755
--- a/forge_key.py
+++ b/forge_key.py
@@ -15,7 +15,7 @@ def main():
     args = arg_parser.parse_args()
 
     # Generate RSA key pair
-    data_len = 31
+    data_len = 61
     data = random.getrandbits(data_len * 8).to_bytes(data_len, 'little')
     key = forge_rsa_key(data)
 
@@ -30,59 +30,44 @@ def main():
         )
         print(response.service_id)
 
+def crypto_prime(num):
+    i = 2
+    for x in range(64): # 2^(-128) false
+        if not gmpy2.is_strong_prp(num, i): return False
+        i = gmpy2.next_prime(i)
+    return True
 
-def forge_rsa_key(data: bytes, key_size=1024, data_size=254):
-    prime_size = key_size // 2                # Size of p and q
-    assert (data_size + 2) <= prime_size      # Reserve MSB and LSB bits
-    assert len(data) * 8 <= data_size         # Data size sanity check
-    data = int.from_bytes(data, 'little')
-
-    # Randomly generate p and q
+def next_prime(num):
     while True:
-        # Pick random p. Set LSB and MSB to 1
-        p = random.getrandbits(prime_size) | (2 ** (prime_size - 1) + 1)
-        if not gmpy2.is_strong_prp(p, 2):
-            continue
-
-        # q suffix, or lower bits of q, are derived from data and p
-        # q prefix, or higher bits of q, are randomly selected
-        q_suffix_size = data_size + 1
-        q_prefix_size = prime_size - q_suffix_size
-
-        m = 2 ** q_suffix_size
-        p_inv = int(gmpy2.invert(p, m))
-        q_suffix = (((data << 1) | 1) * p_inv) % m
-        assert (q_suffix & 1) == 1
-
-        # Generate random q prefix and run prime test
-        found = False
-        for _ in range(1024):
-            q_prefix = random.getrandbits(q_prefix_size) | (2 ** (q_prefix_size - 1))
-            q = (q_prefix << q_suffix_size) | q_suffix
-
-            if gmpy2.is_strong_prp(q, 2):
-                found = True
-                break
-
-        if found:
-            break
+        num = gmpy2.next_prime(num)
+        # ensure a sufficient is-prime confidence
+        if crypto_prime(num): return int(num)
+
+DATA_LEN = 488
+
+def forge_rsa_key(data: bytes, key_size = 1024):
+    # By Cramer's conjecture (which is likely to be true),
+    #  the maximum prime gap for a 512-bit number should be around 10**5.
+    #  Thus, we choose 488-bit for our data size, which allows a maximum
+    #  gap of 2^(511-488) ~ 5*10**6.
+    # There is no need of a nonce, since there is enough randomness in the
+    #  chosen prime factor.
+    assert(len(data) == DATA_LEN // 8)
+    # let the highest bit be 1
+    n_expect = 1 << 1023 | int.from_bytes(data, 'little') << (1023 - DATA_LEN)
+    while True:
+        p = next_prime(random.getrandbits(512))
+        q = next_prime((n_expect - 1) // p + 1)
+        n = p * q
+        # final (paranoid) correctness check,
+        #  should be true given the conjecture is true
+        if (n >> (1023 - DATA_LEN)) == (n_expect >> (1023 - DATA_LEN)): break
+        print('meow')
 
     # Compute RSA components
-    n = p * q
     e = 65537
     d = int(gmpy2.invert(e, (p - 1) * (q - 1)))
 
-    # Test correctness
-    recovered_data = (n >> 1) & (2 ** data_size - 1)
-    assert gmpy2.is_strong_prp(p, 2) and gmpy2.is_strong_prp(q, 2)
-    assert data == recovered_data
-    # print('p', bin(p))
-    # print('q', bin(q))
-    # print('n', bin(p * q))
-    # print('d', bin(data))
-    # print('r', bin(recovered_data))
-    # print('d', bin(data))
-
     # Library reference
     # https://pycryptodome.readthedocs.io/en/latest/src/public_key/rsa.html
     key = RSA.construct(
@@ -96,6 +81,9 @@ def forge_rsa_key(data: bytes, key_size=1024, data_size=254):
     ret = str(base64.b64encode(der), 'ASCII')
     return ret
 
+def data_from_public_key(n):
+    data_num = (n - (1 << 1023)) >> (1023 - DATA_LEN)
+    return data_num.to_bytes(DATA_LEN // 8, 'little')
 
 if __name__ == '__main__':
     main()