the helper fuctions for a url shortener

Question

Currently I am still considering of using the get_password instead of the token_urlsafe but not sure yet. The shortener will be offering also encryption in case someone needs it for private short urls. Do you see something weak in the encryption?

The code:

from secrets import SystemRandom, choice, token_urlsafe
from sqids import Sqids
from random import sample
import nacl.secret
import nacl.utils

my_secure_rng = SystemRandom()
main_alphabet = "123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz"
punct = "-_~"


def get_short_url(srng=my_secure_rng, ab: str = main_alphabet) -> str:
    """Given a secrets.SystemRandom and an alphabet return a random string"""
    id_db = [srng.randrange(100, 9999999999999)]
    rng_ab = "".join(sample(list(ab), k=len(ab)))  # shuffle the alphabet
    sqids = Sqids(min_length=9, alphabet=rng_ab)
    s_url = sqids.encode(id_db)  # random string
    return s_url


def get_password(lg: int = 12, ab: str = (main_alphabet + punct)) -> str:
    """Get a secrets.SystemRandom password"""
    while True:
        password = "".join(choice(ab) for i in range(lg))
        if (
            any(c.islower() for c in password)
            and any(c.isupper() for c in password)
            and sum(c.isdigit() for c in password) >= 2
        ):
            break
    return password


def get_short_encr(encrypt: int, long_url: str):
    """Get a shorturl and a token to decrypt the long url
    (in case you want it encrypted )"""
    if encrypt == 1:
        short_url = get_short_url()
        key = token_urlsafe(24).encode("utf-8")
        box = nacl.secret.Aead(key)
        token = box.encrypt(bytes(long_url, "utf-8"))

    else:
        short_url = get_short_url()
        token = -1  # NOTE: use 0? or something else?
    return {short_url, token}

Answer 1

elided characters

When defining main_alphabet it would be helpful to comment on the goal of wanting just visually distinct glyphs, and describing the ones you found ambiguous. Better than a comment would be to express that in code:

from string import digits, ascii_letters

main_alphabet = "".join(sorted((set(digits + ascii_letters) - set("0OIl"))))

my_secure_rng is an ok identifier, but a shorter rng would have sufficed.

design of public API

The ab in the signature should be spelled alphabet. A short name like that could be OK for a local variable, but here we're exposing it to callers so it should be self explanatory.

Similarly for the lg "length" parameter in get_password().

Consider renaming the first function to _get_short_url, if it is a private helper that you don't expect apps to be calling directly.

human readable number

    id_db = [srng.randrange(100, 9999999999999)]

This would be much more readable as

    id_db = [srng.randrange(100, 9_999_999_999_999)]

or better you could write int(1e13 - 1). A more relevant bound might be 2 ** 43 (8_796_093_022_208), showing that we're sticking below the sqids (signed) limit of 63 bits.

Its unclear why 99 would be a bad ID, and besides the chance of choosing a two-digit number is pretty low. I imagine that "all nine's" would be an acceptable ID, so consider writing it as srng.randrange(0, int(1e13)). Notice that range(a, b) will never yield b, and similarly .randrange(a, b) will never return b.

already a sequence

    rng_ab = "".join(sample(list(ab), k=len(ab)))

Since ab is a str it is already a sequence, so no need to ask list() to produce another sequence.

Consider assigning to a better identifier, perhaps rand_ab, since the resulting string is in no way a random number generator.

The identifier id_db seems to be trying to avoid shadowing the builtin id() function -- good! But db_id would have been a more natural identifier, and perhaps following the "add a suffix" convention of id_ would be the best fit.

The docstring promises to "return a random string", which is somewhat accurate, but only half the story. It's a string that corresponds to a single random integer, which squids .decode() can recover. Also, that integer is range limited, so there's a great many nine-character strings which "look random" yet would not properly decode.

There's a lot of structure in the return string, and such strings can be readily distinguished from e.g. nine rolls of a D58.

It's a stretch to claim this function returns a short URL, as the value doesn't begin with "https://" or similar URL scheme. We do get a perfectly nice URL path.

docstring

The get_password() docstring is accurate. It should be a bit longer, describing the "two out of three character classes" post-condition that it promises.

Consider breaking out the character class test as a predicate helper function. Then you could replace while True: ... break with a more natural while not is_valid(password):

It's odd that punct is not viewed as a character class here. Also, notice that token_urlsafe() outputs two of the three defined punctuation characters: _ and -

type stability

This is a surprising choice:

def get_short_encr(encrypt: int, ... ):

Accepting a bool parameter would make much more sense.

Consider giving this function a longer name, without "encr" abbrev.

It would be helpful to type annotate the return value. Please reconsider what you return. As written, we return one of

• set[str | int]
• set[str | bytes]

It's hard to see how caller would properly make use of this 2-element set. Wouldn't you rather return a tuple? Better, define a namedtuple.

---

Do you see something weak in the encryption?

Hard to say, as you did not write down a Threat Model or Concept of Operations. We don't know what threat actors you're trying to defend against, nor the costs of a breach.

The docstring explains that "AEAD messages are repudiable", which may or may not be a fit for your use case.

You didn't specify security parameters such as NONCE_SIZE, and I can't tell if its default is suitable for you.

We see just the first two lifecycle stages:

1. roll a key
2. encrypt
3. transmit
4. decrypt
5. expire / revoke / re-roll key

In a deployed system, key management is the hard part to get right. You have defined a Public API which weirdly mixes key creation and encryption together. Consider forcing the caller to make separate calls, one focused on the key and another call focused on using the key on some data. We'd expect to see a handful of key calls, and lots of data calls.

short key

Choosing a 96-bit key with token_urlsafe(24) seems adventurous. The docstring recommends using a longer key. If there is some engineering constraint that motivated choosing a deliberately weaker security parameter, it should be written down in the source code or in a Requirements document.

Answer 2

This is the code update after taking into consideration the review, in case someone is interested to know.

import base64
from collections import namedtuple
from nacl import secret
from nacl import utils
from pydantic_core import Url
from pydantic import BaseModel, ValidationError, UrlConstraints
from random import sample
from secrets import SystemRandom, choice
from sqids import Sqids
from typing_extensions import Annotated


my_secure_rng = SystemRandom()
main_alphabet = "123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz"
punct = "-_~"
AnyUrl = Url
HttpUrl = Annotated[
    Url,
    UrlConstraints(max_length=2083, allowed_schemes=["magnet", "http", "https"]),
]


def _get_short_url_path(srng=my_secure_rng, ab: str = main_alphabet) -> str:
    """Given a secrets.SystemRandom and an alphabet return an sqid"""
    db_id = [srng.randrange(100, 8_796_093_022_208)]
    rand_ab = "".join(sample(ab, k=len(ab)))  # shuffle the alphabet
    sqids = Sqids(min_length=9, alphabet=rand_ab)
    s_url = sqids.encode(db_id)  # "random" string
    return s_url


def get_password(lenght: int = 12, ab: str = (main_alphabet + punct)) -> str:
    """Get a secrets.SystemRandom password"""
    while True:
        password = "".join(choice(ab) for i in range(lenght))
        if (
            any(c.islower() for c in password)
            and any(c.isupper() for c in password)
            and sum(c.isdigit() for c in password) >= 2
        ):
            break
    return password


def get_short_path(encrypt: bool, long_url: str) -> (str, str, str):
    """Get a short url path, the (encrypted)
    long url, and a token to decrypt the long url
    (in case you want it encrypted )"""
    ent_url = namedtuple("ENT_URL", ("short_path", "e_url", "key"))
    if encrypt:
        short_path = _get_short_url_path()
        key = utils.random(secret.Aead.KEY_SIZE)
        box = secret.Aead(key)
        e_url = box.encrypt(bytes(long_url, "utf-8"))
        e_url = base64.b64encode(e_url).decode("utf-8")
        key = base64.b64encode(key).decode("utf-8")
    else:
        short_path = _get_short_url_path()
        e_url = "False"  # NOTE: revisit
        key = "False"
    ent = ent_url(short_path, key, e_url)
    return ent


def decrypter(key: str, e_url: str) -> str:
    """Using a Aead key and an encrypted url string
    give back the long url string"""
    box2 = secret.Aead(base64.b64decode(key))
    url = box2.decrypt(base64.b64decode(e_url)).decode("utf-8")
    return url


def raise_bad_request(message):

    raise ValidationError(status_code=400, detail=message)


def is_safe_url(long_url: str) -> bool:
    """Check if the given url is safe and an actual url"""

    # FIX: check urls and scan using some kind of list
    class MyModel(BaseModel):
        url: HttpUrl

    try:
        MyModel(url=long_url)
        a = True
    except ValidationError as e:
        a = False
    return a
```