diff --git a/colorhash/__main__.py b/colorhash/__main__.py index 763ff0e..d1162e2 100644 --- a/colorhash/__main__.py +++ b/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 = [ - # red - [f"#{0x110000 * i:06x}" for i in range(0x10)], - # green - [f"#{0x001100 * i:06x}" for i in range(0x10)], - # blue - [f"#{0x000011 * i:06x}" for i in range(0x10)], - # black - [f"#{0x111111 * 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), +PALETTES = { + "red": [f"#{0x110000 * i:06x}" for i in range(0x10)], + "green": [f"#{0x001100 * i:06x}" for i in range(0x10)], + "blue": [f"#{0x000011 * i:06x}" for i in range(0x10)], + "black": [f"#{0x111111 * 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)], } -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: - inpath = sys.argv[1] - if inpath != "-": - infile = open(inpath, "rb") + ap.add_argument( + "infile", + type=argparse.FileType("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 + ############################################################################ + + # Get the hash + # file_digest (I hope) will not load too much into memory + hashdata = hashlib.file_digest(args.infile, args.hash).digest() + + # Choose the palette + 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 SVG + matrix = matricizer.matricize(hashdata) + colors = colorizer.colorize(matrix) + svg = gensvg(colors, args.square_size) + if str(args.out) == '-': + sys.stdout.write(svg) + else: + args.out.write_text(svg) -hashdata = hashlib.file_digest(infile, hash_algo).digest() -w, h = dimensions_table[hash_algo] - -palette_no = sum(hashdata) % 8 -palette = color_table[palette_no] - -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) -colors = PaletteColorizer(palette).colorize(matrix) -print(gensvg(colors, 32)) +main() diff --git a/colorhash/matricizer.py b/colorhash/matricizer.py index 6713046..92ec1f6 100644 --- a/colorhash/matricizer.py +++ b/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): - pass +class RandomartMatricizer(Matricizer): + 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