pwgen utility -- in Python

pwgen.py

#!/usr/bin/env python3
# SPDX-License-Identifier: MPL-2.0
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at https://mozilla.org/MPL/2.0/.
from __future__ import annotations

import argparse
import random
import secrets
import string
import sys
import time
from typing import Final, Protocol, cast

SIMPLE_PUNCTUATION: Final[str] = "!@#$%_"
MIN_PASS_LEN: Final[int] = 8
MIN_PART_WT: Final[float] = 5.0
MIN_SHUFFLE_CAP: Final[int] = 100


class _Options(Protocol):
    length: int
    num: int
    all_punct: bool
    punct: str


def _get_options() -> _Options:
    parser = argparse.ArgumentParser(
        "pwgen",
        description="Securely generate strong passwords (comprising UPPPERCASE, lowercase, Numbers, and Punctuations)",
    )

    parser.add_argument("-n", "--num", type=int, default=1, help="Number of passwords to generate. Defaults to 1")
    parser.add_argument(
        "--all-punct", action="store_true", default=False, help="If specified, use all punctuations. Overrides --punct"
    )
    parser.add_argument(
        "--punct", metavar="CHARS", default=SIMPLE_PUNCTUATION, help="Characters to include in the 'punctuation' class"
    )
    parser.add_argument("length", type=int, help="Length of each generated password")

    _opts = cast(_Options, parser.parse_args())
    return _opts  # noqa: RET504


def gen_password(
    length: int,
    *,
    punctuations: str = string.punctuation,
    filler: str = string.ascii_lowercase,
    part_weight: float = 6.0,
    shuffle_cap: int = 100,
) -> str:
    """
    Generates in a secure way, a password of a certain length with *guaranteed* characters from the
    UPPERCASE, lowercase, Digits, and Punctuation classes.

    :param length: Length of password, must be >= 8
    :param punctuations: Override the Punctuation class
    :param filler: Character class to fill up the character up to specified length
    :param part_weight: Weight used to generate guaranteed characters. Must be >= 5.0
    :param shuffle_cap: The cap (maximum) number that shuffling will be performed. Must be >= 100
    :return: The securely-generated password
    :raises ValueError: if length < 8, or part_weight < 5.0, or shuffle_cap < 100
    """
    if length < MIN_PASS_LEN:
        raise ValueError("length must be 8 or more")
    if part_weight < MIN_PART_WT:
        raise ValueError("part_weight must be 5.0 or more")
    if shuffle_cap < MIN_SHUFFLE_CAP:
        raise ValueError("shuffle_cap must be 100 or more")
    chars = []

    partlen = max(1, round(length / part_weight))
    for _ in range(partlen):
        chars.append(secrets.choice(string.ascii_lowercase))
        chars.append(secrets.choice(string.ascii_uppercase))
        chars.append(secrets.choice(string.digits))
        chars.append(secrets.choice(punctuations))

    lastlen = length - len(chars)
    chars.extend(secrets.choice(filler) for _ in range(lastlen))

    # Prevent timing attacks by shuffling an unknown number of times, and sleeping in each iteration an unknown
    # number of milliseconds
    for _ in range(10 + secrets.randbelow(shuffle_cap)):
        random.shuffle(chars)
        time.sleep(secrets.randbelow(10) / 1000.0)
    return "".join(chars)


def _main(opts: _Options) -> int:
    if opts.length < MIN_PASS_LEN:
        print(f"ERROR: Password length must be >= {MIN_PASS_LEN}", file=sys.stderr)
        return 1

    if opts.all_punct:
        opts.punct = string.punctuation

    # noinspection PyBroadException
    try:
        for _ in range(opts.num):
            password = gen_password(opts.length, punctuations=opts.punct)
            print(password)
    except Exception as e:
        print(f"ERROR: [{type(e)}]{e}", file=sys.stderr)
        return 2
    else:
        return 0


if __name__ == "__main__":
    sys.exit(_main(_get_options()))