#!/usr/bin/env python3 from sys import stdin from sys import exit as sysexit from io import BytesIO import click from blake3 import blake3 from . import Identicon, get_deterministic_stream, ClosableStream DIGEST_SIZE = 20 BUF_SIZE = 65536 # Linux default pipe capacity is 64KiB (64 * 2^10) @click.command( help=( 'Generate OpenSSH style randomart identicon for arbitrary data.\n\n' 'If TEXT or --file is not supplied, data is read from STDIN.' ) ) @click.argument('text', default=None, type=str, required=False) @click.option( '--file', '-f', default=None, type=click.Path(exists=True), help='Calculate from file or directory (recursive).' ) @click.option( '--fingerprint', '-p', default=False, required=False, is_flag=True, help='Print fingerprint instead of identicon.' ) def main(**kwargs): if not (stream := get_input_stream(kwargs)): print_usage_and_exit() digest = get_digest(stream.stream) stream.close() if not kwargs.get('fingerprint'): i = Identicon(digest) i.calculate() print(i) else: print(digest.hex()) def get_input_stream(kwargs): stream = None if (text := kwargs['text']) is not None: stream = ClosableStream(BytesIO(text.encode())) elif file := kwargs['file']: stream = get_deterministic_stream(file, BUF_SIZE) elif not stdin.isatty(): stream = ClosableStream(stdin.buffer) return stream def print_usage_and_exit(): command = main with click.Context(command) as ctx: click.echo(command.get_help(ctx)) sysexit(1) def get_digest(stream): # pylint: disable=not-callable hasher = blake3() while data := stream.read(BUF_SIZE): hasher.update(data) return hasher.digest(length=DIGEST_SIZE) if __name__ == '__main__': main()