Add command line parsing

Signed-off-by: Alek Ratzloff <alekratz@gmail.com>

colorhash/__main__.py

@@ -1,69 +1,151 @@
+import argparse
 import hashlib
+from pathlib import Path
 import sys
 
 from .colorizer import PaletteColorizer
-from .matricizer import NibbleMatricizer
+from .matricizer import NibbleMatricizer, RandomartMatricizer
 from .svg import gensvg
 
 
 # TODO - WASM compile for embedding directly in HTML
-# TODO - command line parsing for hash type, infile, forcing a palette, etc
-# TODO - file streaming for infile so we aren't loading e.g. 4GB into memory unnecessarily
 # TODO - option to add a caption based on the filename
 # TODO - palettes defined by JSON
 
 
-color_table = [
+PALETTES = {
-    # red
+    "red": [f"#{0x110000 * i:06x}" for i in range(0x10)],
-    [f"#{0x110000 * i:06x}" for i in range(0x10)],
+    "green": [f"#{0x001100 * i:06x}" for i in range(0x10)],
-    # green
+    "blue": [f"#{0x000011 * i:06x}" for i in range(0x10)],
-    [f"#{0x001100 * i:06x}" for i in range(0x10)],
+    "black": [f"#{0x111111 * i:06x}" for i in range(0x10)],
-    # blue
+    "cyan": [f"#{0x001111 * i:06x}" for i in range(0x10)],
-    [f"#{0x000011 * i:06x}" for i in range(0x10)],
+    "yellow": [f"#{0x111100 * i:06x}" for i in range(0x10)],
-    # black
+    "magenta": [f"#{0x110011 * i:06x}" for i in range(0x10)],
-    [f"#{0x111111 * i:06x}" for i in range(0x10)],
+    "white": [f"#{0x111111 * (0xF - i):06x}" for i in range(0x10)],
-    # cyan
-    [f"#{0x001111 * i:06x}" for i in range(0x10)],
-    # yellow
-    [f"#{0x111100 * i:06x}" for i in range(0x10)],
-    # magenta
-    [f"#{0x110011 * i:06x}" for i in range(0x10)],
-    # white
-    [f"#{0x111111 * (0xF - i):06x}" for i in range(0x10)],
-]
-
-dimensions_table = {
-    "md5": (8, 4),
-    "sha1": (8, 5),
-    "sha224": (8, 7),
-    "sha256": (8, 8),
-    "sha384": (12, 8),
-    "sha512": (16, 8),
 }
 
 
-hash_algo = "sha512"
+def main():
+    MATRIX_CHOICES = {
+        "nibble": "Use each nibble (4 bits) of the hash to generate a matrix",
+        "randomart": "Use the SSH 'randomart' algorithm to generate a matrix",
+    }
+    MATRIX_HELP = "MATRIX STRATEGY (-m, --matrix)\n" + "\n".join(
+        f"    {choice} - {desc}" for choice, desc in MATRIX_CHOICES.items()
+    )
+    PALETTE_CHOICES = [
+        "auto",
+        "red",
+        "green",
+        "blue",
+        "black",
+        "cyan",
+        "yellow",
+        "magenta",
+        "white",
+    ]
+    PALETTE_HELP = "\n".join(
+        [
+            "PALETTE CHOICES",
+            "    " + ", ".join(PALETTE_CHOICES),
+        ]
+    )
+    HASH_CHOICES = ["md5", "sha1", "sha224", "sha256", "sha384", "sha512"]
+    EPILOGUE = "\n\n".join([MATRIX_HELP, PALETTE_HELP])
 
-infile = sys.stdin.buffer
+    ap = argparse.ArgumentParser(
+        description="Create a piece of art based on the hash of a file.",
+        epilog=EPILOGUE,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
 
-if len(sys.argv) > 1:
+    ap.add_argument(
-    inpath = sys.argv[1]
+        "infile",
-    if inpath != "-":
+        type=argparse.FileType("rb"),
-        infile = open(inpath, "rb")
+        default=sys.stdin,
+        help="The input file to use. Set to '-' or blank for STDIN. default: STDIN",
+    )
+    ap.add_argument(
+        "-o",
+        "--out",
+        metavar="OUTFILE",
+        type=Path,
+        default="-",
+        help="The output file to use. Set to '-' or blank for STDOUT. default: STDOUT",
+    )
+    ap.add_argument(
+        "-m",
+        "--matrix",
+        metavar="MATRIX",
+        choices=MATRIX_CHOICES.keys(),
+        default="nibble",
+        help="Choose the strategy that turns the hash into a matrix. default: nibble",
+    )
+    ap.add_argument(
+        "-p",
+        "--palette",
+        metavar="PALETTE",
+        choices=PALETTE_CHOICES,
+        default="auto",
+        help="Choose the palette. default: auto",
+    )
+    ap.add_argument(
+        "-a",  # the "a" is for "algorithm" (since -h is taken)
+        "--hash",
+        metavar="ALGORITHM",
+        choices=HASH_CHOICES,
+        default="sha512",
+        help="Choose the hash algorithm. default: sha512",
+    )
+    ap.add_argument(
+        "-z",
+        "--square-size",
+        metavar="PX",
+        type=int,
+        default=32,
+        help="Decide how big the output squares are, in pixels. default: 32",
+    )
+    args = ap.parse_args()
 
+    ############################################################################
+    # End arg parsing
+    ############################################################################
 
-hashdata = hashlib.file_digest(infile, hash_algo).digest()
+    # Get the hash
-w, h = dimensions_table[hash_algo]
+    # file_digest (I hope) will not load too much into memory
+    hashdata = hashlib.file_digest(args.infile, args.hash).digest()
 
-palette_no = sum(hashdata) % 8
+    # Choose the palette
-palette = color_table[palette_no]
+    palette: list[str]
+    if args.palette == 'auto':
+        palette = list(PALETTES.values())[sum(hashdata) % 8]
+    else:
+        palette = PALETTES[args.palette]
 
+    # Choose the dimensions and the matricizer
+    matricizer: Matricizer
+    match args.matrix:
+        case 'nibble':
+            w, h = NibbleMatricizer.DIMENSIONS[args.hash]
+            matricizer = NibbleMatricizer(w, h)
+        case 'randomart':
+            # 17x9 is what openssh uses
+            # TODO - allow configuring dimensions, maybe
+            matricizer = RandomartMatricizer(17, 9)
+        case _:
+            assert False, f"invalid args.matrix: {args.matrix}"
+
+    # Choose the colorizer
     colorizer = PaletteColorizer(palette)
 
-# Print colors
-# pprint.pprint([[hex(c) for c in row] for row in colors])
-
     # Print SVG
-matrix = NibbleMatricizer(w, h).matricize(hashdata)
+    matrix = matricizer.matricize(hashdata)
-colors = PaletteColorizer(palette).colorize(matrix)
+    colors = colorizer.colorize(matrix)
-print(gensvg(colors, 32))
+    svg = gensvg(colors, args.square_size)
+    if str(args.out) == '-':
+        sys.stdout.write(svg)
+    else:
+        args.out.write_text(svg)
+
+
+main()

colorhash/matricizer.py

@@ -18,6 +18,14 @@ class Matricizer(metaclass=abc.ABCMeta):
 
 
 class NibbleMatricizer(Matricizer):
+    DIMENSIONS = {
+        "md5": (8, 4),
+        "sha1": (8, 5),
+        "sha224": (8, 7),
+        "sha256": (8, 8),
+        "sha384": (12, 8),
+        "sha512": (16, 8),
+    }
     def matricize(self, data: bytes) -> Matrix:
         """
         Convert a set of bytes to a list of rows of nibbles.
@@ -45,5 +53,27 @@ class NibbleMatricizer(Matricizer):
         return cols
 
 
-class PerlinMatricizer(Matricizer):
+class RandomartMatricizer(Matricizer):
-    pass
+    def matricize(self, data: bytes) -> Matrix:
+        """
+        Create a matrix based on the "randomart" algorithm from ssh-keygen.
+        """
+        rows = [[0] * self.w for _ in range(self.h)]
+        c = self.w // 2
+        r = self.h // 2
+        for value in data:
+            for _ in range(4):
+                if value & 0x1:
+                    c += 1
+                else:
+                    c -= 1
+                if value & 0x2:
+                    r += 1
+                else:
+                    r -= 1
+                c = min(max(c, 0), self.w - 1)
+                r = min(max(r, 0), self.h - 1)
+                # max value is 0xf
+                if rows[r][c] < 0xf:
+                    rows[r][c] += 1
+        return rows