How to architect a solution for generating and expiring short codes in a web application?

Question

I'm working on a USSD application where I have a specific requirement to generate short codes for users after they receive a call with a pharmacy practitioner. These short codes should expire after 72 hours or become invalid once used by the user. Additionally, the short codes are used for payment processing, where the system needs to determine the corresponding payment amount based on the short code provided. To minimize user input errors, the short codes will consist of either 4 or 5 digits.

Here's my current approach:

• I'm using quadratic residues to generate random unique codes of 4 digits.
• Instead of creating an additional table or adding the codes as attributes to existing tables, I store the codes and associated user objects in a Redis database, with an expiration time set.
• To ensure uniqueness, I verify if the code is already assigned to another user by checking its existence in the Redis database before assigning it.
• Once the code is sent to the user, it is removed from the Redis database once it has been used.

I'm seeking suggestions and improvements for this architecture. Currently, my stack includes Flask, MySQL, SQLAlchemy, Celery for background tasks, and RabbitMQ as the queuing system.

Any advice or alternative approaches to handle the generation, expiration, and removal of these short codes in an efficient and scalable manner would be greatly appreciated. Thank you!

Here is a code snippet that generates a random 4-digit number with unique digits

import random
import math



def is_prime(n):
    """Checks if a number is prime."""
    if n < 2:
        return False
    for i in range(2, int(math.sqrt(n)) + 1):
        if n % i == 0:
            return False
    return True

def generate_random_prime():
    """Generates a random prime number between 1000 and 9999."""
    while True:
        prime_number = random.randint(1000, 9999)
        if is_prime(prime_number):
            return prime_number

def generate_quadratic_residue(prime_number):
    """Generates a random quadratic residue modulo a prime number."""
    while True:
        quadratic_residue = random.randint(1, prime_number - 1)
        if (quadratic_residue ** 2 % prime_number != 1):
            return quadratic_residue

def generate_unique_code():
    """Generates a unique 4-digit code using quadratic residues."""
    prime = generate_random_prime()
    quadratic_residue = generate_quadratic_residue(prime)
    code = str(quadratic_residue).zfill(4)
    return code

Answer 1

In general, it is not considered good practice to write your own code for things that are supposed to be cryptographically secure, because of many reasons (for example, a lack of auditing/maintenance).

I would suggest using python's secrets, like so:

import secrets

secure_random_number = secrets.randbelow(10000)
...

Here is the documentation.

It will also be much faster than your current implementation (which draws stuff from a haystack one at a time, hoping it will be a needle-prime).

I don't really have any suggestions for the rest of your dataflow, but probably you should encrypt as much as possible the communication between the parts (one example, store a hash and not the actual 4-digit nuber, and then check user input against it after hashing it).

Of course, in these applications scalability is a huge factor. Since the code is kept for 3 days, I am assuming there aren't very many users, but in any case you should see if those things pose additional overhead.