identicon/identicon/__main__.py

#!/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()