RSA encryption and decryption: performance, security, and best practises

Question

I'm working on a RSA encryption and decryption project in python. I've written the following code, but I'm looking for suggest on how to improve it. Specifically, I'm interested in optimizing performance, enhancing security, and following best practise. Any feedback would be greatly appreciated!

import random
import math

# n= p.q
# phi(n) = (p-1)(q-1)

# Prime number checking function
def is_prime(number):
    if number < 2:
        return False
    for i in range(2, int(math.sqrt(number)) + 1):  # Optimized to check up to sqrt(number)
        if number % i == 0:
            return False
    return True

# Generate a random prime between a range
def generate_prime(min_value, max_value):
    prime = random.randint(min_value, max_value)
    while not is_prime(prime):
        prime = random.randint(min_value, max_value)
    return prime

# Modular inverse using the Extended Euclidean Algorithm
def mod_inverse(e, phi):
    def egcd(a, b):
        if a == 0:
            return (b, 0, 1)
        else:
            g, y, x = egcd(b % a, a)
            return (g, x - (b // a) * y, y)
    
    g, x, y = egcd(e, phi)
    if g != 1:
        raise ValueError("Modular inverse does not exist")
    else:
        return x % phi

# Function to encrypt message in chunks
def encrypt_message(message, e, n, block_size):
    message_encoded = [ord(c) for c in message]
    chunks = [message_encoded[i:i + block_size] for i in range(0, len(message_encoded), block_size)]
    ciphertext = []
    
    for chunk in chunks:
        # Convert chunk to an integer (ASCII representation)
        m = int(''.join([f'{c:03d}' for c in chunk]))
        # Encrypt using (m ^ e) mod n
        ciphertext.append(pow(m, e, n))
    
    return ciphertext

# Function to decrypt message in chunks
def decrypt_message(ciphertext, d, n, block_size):
    message_decoded = []
    
    for cipher_chunk in ciphertext:
        # Decrypt using (c ^ d) mod n
        m = pow(cipher_chunk, d, n)
        # Convert back to individual characters from the integer
        chunk_str = str(m).zfill(block_size * 3)  # Zfill ensures we get proper block size
        message_decoded.extend([int(chunk_str[i:i + 3]) for i in range(0, len(chunk_str), 3)])
    
    return ''.join([chr(c) for c in message_decoded])

# RSA key generation
p, q = generate_prime(1000, 50000), generate_prime(1000, 50000)
while p == q:
    q = generate_prime(1000, 50000)

n = p * q
phi_n = (p - 1) * (q - 1)

e = random.randint(3, phi_n - 1)
while math.gcd(e, phi_n) != 1:
    e = random.randint(3, phi_n - 1)

d = mod_inverse(e, phi_n)

print("Prime number P: ", p)
print("Prime number Q: ", q)
print("Public Key (e, n): ", (e, n))
print("Private Key (d, n): ", (d, n))
print("n: ", n)
print("Phi of n: ", phi_n)

# Example message
message = "Khoa Toan-Co-Tin hoc, Truong Dai hoc Khoa hoc Tu nhien"
block_size = len(str(n)) // 3 - 1  # Block size in terms of number of characters (each char takes 3 digits in ASCII)

# Encrypt message
ciphertext = encrypt_message(message, e, n, block_size)
print("Ciphertext: ", ciphertext)

# Decrypt message
decrypted_message = decrypt_message(ciphertext, d, n, block_size)
print("Decrypted Message: ", decrypted_message)

What improvements can I make to this code?

Answer 1

I am not a crypto expert but the random module is discouraged for security purposes:

Warning

The pseudo-random generators of this module should not be used for security purposes. For security or cryptographic uses, see the secrets module.

Source

You could at least use random.seed to add entropy but this is still nowhere near production-grade crypto. There is a weakness in terms of implementation here, because pseudo-random/biased sequences make the algo much less robust.

So the secrets module could be used to generate random numbers in a less deterministic way. Note that this actually uses the random module but a different class: SystemRandom, using OS sources like /dev/urandom. Actually, there is specialized hardware to generate random numbers, and we even have HSM because "true" randomness is something that regular computers can't easily provide.

Since we need "two primes that are large and random, and ideally approximately the same size", the next question is the suitability of your primes. If we are aiming for primes of say, 2048 bit long in size which should be a fair compromise between security and performance, then you'd need significantly bigger numbers. An example should look like this.

Larger numbers will take longer but there are ways to compute in parallel. See for example How to Check if Numbers are Prime Concurrently in Python.

So it looks to me that your numbers can easily be factorized using commodity hardware. I can't speak with authority on the matter but this code presents obvious limitations. For more in depth discussion you might also want to have a look at the other forum I have already linked to: https://crypto.stackexchange.com/

Rolling out home-brew crypto is never a good idea, so the "best" practice would be to rely on a proven implementation. There are specialized libs like Crypto in Python for the purpose of generating RSA keys or prime numbers. In other words, treat this as an exercise but don´t use it in production. Ever.

Answer 2

Cryptography

The first rule of cryptography is don't write cryptographic software. It's fine as an exercise, but not for practical use.

With that out of the way, there's a glaring problem with your algorithm. RSA is predicated on the fact that you can't factorize n in any reasonable amount of time. The general number field sieve is a factorization algorithm running in sub-exponential time in log n, which is better than n^c for any positive c. This is a problem when you look at how you generate your primes.

By the prime number theorem, the probability of a number uniformly picked from [1, N] being a prime is ~1/ln(N), so the expected number of attempts before picking a prime is ~ln(N). For each attempt, you run test divisions as the primality test with run time O(sqrt(N)). This results in a runtime of O(ln(N) * sqrt(N)). Since n ~ N^2, your prime generation takes O(ln(n) * n^(1/4)) time, which is worse than the general number field sieve.

In other words, your current algorithm is slower to execute than an attack. Instead, you need a faster primality test such as those suggested here. It should go without saying, but you will need a much larger n for RSA to make much sense, a common choice being on the order of 2^2048.

---

Performance

This is not how you deal with performance. If you care about performance, you have to benchmark. If you don't have a benchmark, you don't really care about performance.

Additionally, since there are many security problems (such as those mentioned by user555045), the program should be considered incorrect: it doesn't fulfill the purpose of providing security. Optimizing an incorrect program isn't very fruitful.

But if we assume those issues are fixed, optimizing python code isn't very fruitful to begin with. Using a JIT such as numba will likely outperform anything you can do with plain python. If that isn't sufficient, you use another language.

---

Code Style

Comments are meant to explain the non-obvious. Code can express meaning by itself and comments reiterating the code is unnecessary. For example

# Function to decrypt message in chunks
def decrypt_message(ciphertext, d, n, block_size):

This comment is what you'd be able to guess just by looking at the function name and signature. Instead, you can write

def decrypt_message(ciphertext: list[int], d: int, n: int, block_size: int) -> str:
    """Decrypts ciphertext with private key d and public key n.
    ciphertext must be encrypted as blocks of block_size.

    Args:
        ... document parameters ...

    Returns:
        ... document return value ...
    """

A docstring is the standard way of documenting what functions do. This is heavily relied on by various documentation tools. You should choose one and follow its format. Additionally, you should always write type annotations.

Answer 3

This code looks to be for educational purposes, which is fine.. It will always have a timing side-channel, Python BigInts are not constant-time.

# Prime number checking function
def is_prime(number):
    if number < 2:
        return False
    for i in range(2, int(math.sqrt(number)) + 1):  # Optimized to check up to sqrt(number)
        if number % i == 0:
            return False
    return True

This primality testing algorithm does not scale to the size of numbers that you need for RSA to be even a little bit secure. A number that has two factors below 50000 can be factored by hand even, and is trivial for a computer. 512-bit keys can be factored by regular people these days, so 256-bit p and q are too small. A 256-bit number has an 128-bit square root, 2^128 iterations is just too much (and this is for numbers that are too small). Sure you can reduce that number a bit, and you can test a handful of numbers in parallel, all of that is in the wrong order of magnitude - it's not enough to make this a thousand time faster, it has to be done a billion billions times faster in order to make it possible at all and even then it would be slow.

Realistically you'd use something like Baillie–PSW, or several rounds of Miller–Rabin optionally combined with a Lucas test, something like that. These don't guarantee weeding out all composite numbers, but you can make the probability of a false negative arbitrarily low, and most critically: these tests are feasible to do on numbers of the size that you need to make an RSA key that isn't trivial to factor.

There are also some people who say you should avoid numbers p (for both p and q of course) such that p - 1 or p + 1 (and some extra things that add up to: maybe p and q should be chosen to be strong primes) are smooth (made up of low factors) to avoid eg Pollard's p - 1 factorization algorithm, but also some (including Ronald Rivest, the R from RSA) who say that doesn't really help because elliptic curve factorization exists which has a similar effect but doesn't require p - 1 or p + 1 to be smooth, I'm not picking a side in that discussion I'm just telling you it exists.

Also if p and q are too close, their product is vulnerable to Fermat's factorization method.

Another way you get into "mathematical trouble" is if e and m are small enough that m^e < n and then the message can be decoded without the private key by taking the e- th root. There are more bad cases, many of which can be avoided by employing some padding scheme (which is fairly subtle and numerous mathematical vulnerabilities have been found with them too).

# Modular inverse using the Extended Euclidean Algorithm

That's a way to do it. Since you only use it for mod_inverse(e, phi_n) and by definition gcd(e, phi_n) = 1, you can also use pow(e, -1, phi_n).

m = int(''.join([f'{c:03d}' for c in chunk]))

I think this is an odd way to do it, I would reinterpret each chunk (but with padding) as an integer directly, without the extraneous string manipulation. And if I was going to concatenate things this way, I would do it in hexadecimal, not decimal.

Answer 4

Your custom function mod_inverse(e, phi) can be replaced by pow(e, -1, phi) since Python 3.8.

Your function generate_prime can be simplified to:

def generate_prime(min_value, max_value):
    while True:
        prime = random.randint(min_value, max_value)
        if is_prime(prime):
            return prime

In cases like the following, you can replace the list comprehension with a generator comprehension: ''.join([f'{c:03d}' for c in chunk]) --> ''.join(f'{c:03d}' for c in chunk) (removed the square brackets).

Answer 5

Note that this is fine for playing around with RSA, but:

1. RSA is very hard to protect against side channel attacks. Other applications running (on the same core) may e.g. try and find the private exponent by looking at CPU usage.
2. RSA requires padding such as OAEP padding to be secure.
3. RSA isn't efficient when encrypting large data sizes, your use a hybrid cryptosystem instead of encrypting chunks. For one it should expand the ciphertext (due to the required padding, see the previous remark). Decryption will be rather slow due to the size of the private exponent (d).
4. Modern cryptography operates on bits and bytes, so you'd encode your message to UTF-8 or similar instead of using ord to turn it into integers.

So please only use the code for learning purposes. Use a library for anything else.

The side channel issue was already covered by another answer, but I've repeated it here because side channel attacks are such a big issue, even in existing RSA implementations. Quite a few of them - including hardware based implementations - do explicitly not protect the key pair generation part of RSA, as it is extremely hard to fully secure.