Add command line parsing

Signed-off-by: Alek Ratzloff <alekratz@gmail.com>
This commit is contained in:
2024-05-22 10:45:41 -07:00
parent f2627f5096
commit 962d2ad316
2 changed files with 164 additions and 52 deletions

View File

@@ -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
############################################################################
hashdata = hashlib.file_digest(infile, hash_algo).digest()
w, h = dimensions_table[hash_algo]
# Get the hash
# 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
palette = color_table[palette_no]
# 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 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))
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)
main()

View File

@@ -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