random_prime.py

from os import urandom
from functools import reduce
from gmpy2 import gcd, mpz, xmpz, is_strong_prp, is_strong_selfridge_prp

SMALL_PRIMES = [ # the smallest 100 primes
      2,   3,   5,   7,  11,  13,  17,  19,  23,  29,
     31,  37,  41,  43,  47,  53,  59,  61,  67,  71,
     73,  79,  83,  89,  97, 101, 103, 107, 109, 113,
    127, 131, 137, 139, 149, 151, 157, 163, 167, 173,
    179, 181, 191, 193, 197, 199, 211, 223, 227, 229,
    233, 239, 241, 251, 257, 263, 269, 271, 277, 281,
    283, 293, 307, 311, 313, 317, 331, 337, 347, 349,
    353, 359, 367, 373, 379, 383, 389, 397, 401, 409,
    419, 421, 431, 433, 439, 443, 449, 457, 461, 463,
    467, 479, 487, 491, 499, 503, 509, 521, 523, 541,
]
SMALL_PRIMES_PRODUCT = reduce(lambda a, b: a*b, map(mpz, SMALL_PRIMES))


def random_prime(bits: int) -> mpz:
    q, r = divmod(bits, 8)
    initial_bytes = bytearray(urandom(q + (r > 0)))

    if r:
        initial_bytes[0] &= (1 << r) - 1   # clear top bits
        initial_bytes[0] |= (1 << (r - 1)) # ensure highest bit is set
    else:
        initial_bytes[0] |= 128

    initial_bytes[-1] |= 1 # ensure lowest bit is set

    p = xmpz(int.from_bytes(initial_bytes, 'big', signed=False))
    pmax = (1 << bits)
    while p < pmax:
        if gcd(p, SMALL_PRIMES_PRODUCT) > 1:
            p += 2
            continue

        if not all([is_strong_prp(p, m) for m in SMALL_PRIMES]):
            p += 2
            continue

        if not is_strong_selfridge_prp(p):
            p += 2
            continue

        return p.make_mpz()

    return mpz(-1)