Generates passwords and PINs.

generate_secrets.py

#!/usr/bin/env python

from enum import Enum
import json
import secrets
import string
import sys
from typing import Iterable, Optional


class CharacterClass(Enum):
    digit = string.digits
    lower = string.ascii_lowercase
    upper = string.ascii_uppercase
    special = "_"


class DEFAULTS:
    length = 16
    must_include = {CharacterClass.digit, CharacterClass.lower, CharacterClass.upper}


def generate(
    length: Optional[int] = None,
    *,
    must_include: Iterable[CharacterClass] = DEFAULTS.must_include,
    human_friendly: bool = True,
    special_chars: Optional[str] = None,
) -> str:
    character_classes = {class_: class_.value for class_ in CharacterClass}

    if not must_include:
        classes = ", ".join(f"'{class_.name}'" for class_ in CharacterClass)
        raise ValueError(f"must_include must include one of {classes}")

    if special_chars is not None:
        if not special_chars:
            raise ValueError("special_chars cannot be empty")
        character_classes[CharacterClass.special] = special_chars

    alphabet = "".join(character_classes[class_] for class_ in must_include)

    if human_friendly:
        alphabet = "".join(char for char in alphabet if char not in "01IOl")

    pw_length = length if length is not None else DEFAULTS.length

    def generate() -> str:
        return "".join(secrets.choice(alphabet) for _ in range(pw_length))

    def valid(password: str) -> bool:
        for class_ in must_include:
            if not any(char in character_classes[class_] for char in password):
                return False
        return True

    while not valid(password := generate()):
        pass

    return password


def passphrase(length: int = 5) -> str:
    with open("dictionary.json") as f:
        words = list(json.load(f).keys())
    return " ".join(secrets.choice(words) for _ in range(length))


def pin(length: int = 4) -> str:
    return str(secrets.randbelow(10**length)).zfill(length)


if __name__ == "__main__":
    length: Optional[int]
    try:
        length = int(sys.argv[1])
    except IndexError:
        length = None
    except ValueError:
        raise ValueError("first argument must be an integer")  # , file=sys.stderr)
        # sys.exit(2)
    if length is not None and length < 1:
        print("first argument must be a positive integer", file=sys.stderr)
        sys.exit(1)

    must_include = DEFAULTS.must_include

    def parse_boolean(arg: str) -> bool:
        arg = arg.lower()
        if arg in ["true", "t", "1"]:
            return True
        if arg in ["false", "f", "0"]:
            return False
        raise ValueError("not a boolean")

    try:
        if parse_boolean(sys.argv[2]):
            must_include |= {CharacterClass.special}
    except IndexError:
        pass
    except ValueError:
        print("second argument must be a boolean", file=sys.stderr)
        sys.exit(2)

    print(generate(length, must_include=must_include))