wordmap/wordmap.py

#!/usr/bin/env python3

from sys import argv
from hashlib import blake2b


def hashfun(input: str, digest_size: int = 4):  # digest size is in bytes
    hash = blake2b(bytes(input, "utf-8"), digest_size=digest_size)
    input_hash = int(hash.hexdigest(), 16)
    return input_hash


def load_words(wordlist_file: str):
    with open(wordlist_file) as word_file:
        valid_words = sorted(set(word_file.read().split("\n")))
    return valid_words


def scale(val: int, src, dst):
    # Scale the given value from the scale of src to the scale of dst.
    return ((val - src[0]) / (src[1]-src[0])) * (dst[1]-dst[0]) + dst[0]


def wordmap(wordlist_file: str, input: str, digest_size: int = 4):
    words = load_words(wordlist_file)
    words_max = len(words)
    digest_max = int.from_bytes(b'\xff' * digest_size, "big")
    input_hash = hashfun(input, digest_size=digest_size)
    # digest_max+1 is so that it can round up to the max
    input_hash_mapped = scale(input_hash, (0, digest_max+1), (0, words_max))
    # alternatively, change round() to math.floor()
    input_hash_mapped = round(input_hash_mapped)
    print(words[input_hash_mapped], end="")


if __name__ == "__main__":
    def usage():
        print(
            f"Usage: {argv[0]} " +
            "<wordlist file> " +
            "<string to hash> " +
            "[string to append]")
        quit(1)

    try:
        if argv[1] == "-h" or argv[1] == "--help":
            usage()
        wordlist_file = argv[1]
        input = argv[2]
        try:
            extrawords = argv[3]
        except IndexError:
            extrawords = ""
    except IndexError:
        usage()
        # input = "example"
        # wordlist_file = "wordlists/term16color.txt"
    wordmap(wordlist_file, input)
    print(extrawords+"\x1b[0m")