pgen.py

#!/usr/bin/env python3

from getpass import getpass
import argparse
import base64
import hashlib
import json
import logging
import os
import secrets


SHORT_SALT_BYTES = 10000


def parse_args():
    parser = argparse.ArgumentParser(description="Generate a password")
    parser.add_argument("domains", nargs="+")
    parser.add_argument("--hash", default="sha512")
    parser.add_argument("--salt", dest="salt_filename", default="~/.pgen.salt")
    parser.add_argument("--checksum", dest="checksums_filename", default="~/.pgen.checksums")
    parser.add_argument("--encoding", default="b85encode")
    parser.add_argument("--add", action="store_true")
    parser.add_argument("--pepper", action="store_true", help="Domain-specific salt")
    parser.add_argument("--verbose", "-v", action="store_true")
    return parser.parse_args()


def read_salt(salt_filename):
    logging.debug(f"Reading salt from {salt_filename}")
    try:
        salt = open(os.path.expanduser(salt_filename), mode="rb").read()
    except FileNotFoundError:
        logging.warning(f"Generating new salt into {salt_filename}")
        salt = secrets.token_bytes(1024 * 1024)
        open(os.path.expanduser(salt_filename), mode="wb").write(salt)
    if len(salt) < SHORT_SALT_BYTES * 2:
        raise Exception("Salt is unusually short")
    logging.debug(f"Salt signature: {base64.b85encode(hashlib.sha512(salt).digest()).decode()}")
    return salt


def read_checksums(checksums_filename):
    logging.debug(f"Reading checksums from {checksums_filename}")
    try:
        checksums = json.loads(open(os.path.expanduser(checksums_filename), mode="rb").read())
    except FileNotFoundError:
        logging.warning(f"Checksums file {checksums_filename} not found")
        checksums = {}
    logging.debug(f"Checksums known: {len(checksums)}")
    return checksums


def write_checksums(checksums_filename, checksums):
    logging.debug(f"Writing checksums to {checksums_filename}")
    with open(os.path.expanduser(checksums_filename), mode="w") as checksums_file:
        checksums_file.write(json.dumps(checksums, indent=4, sort_keys=True))


def get_digest(args, *digest_args):
    prehash = b"".join(
        arg.encode() if isinstance(arg, str) else arg
        for arg in list(digest_args)
    )
    if len(prehash) < SHORT_SALT_BYTES:
        raise Exception("Digest args are unusually short")
    return getattr(hashlib, args.hash)(prehash).digest()


def get_checksum(args, domain, shortpass, salt):
    digest = get_digest(args, domain, shortpass, salt)
    checksum = base64.b85encode(digest).decode()[:20]
    logging.debug(f"{domain}: Checksum computed: {checksum}")
    return checksum


def get_longpass(args, domain, shortpass, salt, config):
    digest = get_digest(args, domain, shortpass, salt, config.get("pepper", ""))
    encoding = getattr(base64, config["encoding"])(digest).decode()
    return "".join((
        config.get("prefix", ""),
        encoding[:config.get("length", 20)],
        config.get("suffix", ""),
    ))


def handle_domain(domain, shortpass, salt, checksums, args):
    shortsalt, longsalt = salt[:SHORT_SALT_BYTES], salt[SHORT_SALT_BYTES:]
    checksum = get_checksum(args, domain, shortpass, shortsalt)
    config = checksums.get(checksum)

    if args.add:
        if config:
            logging.warning(f"{domain}: Checksum already present")
        else:
            config = {
                "encoding": args.encoding,
                "length": 20,
            }
            checksums[checksum] = config
            logging.info(f"{domain}: Checksum added")

    if not config:
        raise Exception(f"{domain}: Checksum not found")

    if args.pepper:
        config["pepper"] = base64.b64encode(secrets.token_bytes(4))[:4].decode()

    return get_longpass(args, domain, shortpass, longsalt, config)


def display_results(results):
    width = max(map(len, results))
    for domain, longpass in results.items():
        print(f"""{domain:>{width}}: {longpass}""")


if __name__ == "__main__":
    args = parse_args()
    logging.basicConfig(
        format="%(levelname)9s: %(message)s",
        level=logging.DEBUG if args.verbose else logging.INFO,
    )
    try:
        shortpass = getpass("Password? ")
        if args.add:
            assert shortpass == getpass("Password (confirm)? ")
        salt = read_salt(args.salt_filename)
        checksums = read_checksums(args.checksums_filename)

        results = {}
        for domain in sorted(args.domains):
            try:
                results[domain] = handle_domain(domain, shortpass, salt, checksums, args)
            except Exception as e:
                logging.error(e)
        if results:
            display_results(results)

        if args.add or args.pepper and results:
            write_checksums(args.checksums_filename, checksums)
    except KeyboardInterrupt:
        print()

The checksum file has been augmented to save configurations for each domain/password combination. This allows:

• better control in cases where domains have unreasonable password restrictions (length, characters, etc)
• per-checksum salts (called "peppers" in the script) to accommodate domains that force password changes at regular intervals