from hashlib import new
from sys import argv

'''
    Copyleft (c) 2021 k98kurz

    Permission to use, copy, modify, and/or distribute this software for any
    purpose with or without fee is hereby granted, provided that the above
    copyleft notice and this permission notice appear in all copies.

    THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
    WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
    MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
    SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
    WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION
    OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
    CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
'''

# returns the calculated difficulty
def calculate(hash: bytes, bigendian: bool = True) -> int:
    # prepare variables
    number = int.from_bytes(hash, byteorder = 'big' if bigendian else 'little')
    length = len(hash)
    trailing_bits = 0

    # shift off bits until a null result is achieved
    while number != 0 and trailing_bits < length * 8:
        number = number >> 1
        trailing_bits += 1

    # difficulty achieved is the preceeding null bits
    return length * 8 - trailing_bits


# checks if a hash has the correct difficulty (preceeding null bits)
def check(diff: int, hash: bytes, bigendian: bool = True) -> bool:
    # prepare variables
    number = int.from_bytes(hash, byteorder='big' if bigendian else 'little')
    length = len(hash)

    # raise exception if difficulty is impossible to achieve
    if diff > length * 8:
        raise Exception('Difficulty cannot exceed bit length of input')

    # shift off all but {diff} preceding bits
    preceding_bits = number >> (length * 8 - diff)

    # it meets the difficulty if preceding bits are 0
    return preceding_bits == 0


# cli
def main(args):
    op = args[0].split(':')
    mode = args[1].split(':')
    input = args[2]
    bigendian = True if len(args) < 4 else bool(args[3])

    if len(op) == 1:
        # default op, e.g. calculate
        if len(mode) == 1:
            # default mode, e.g. no hashing
            input = bytearray.fromhex(input) if mode[0] == 'hex' else bytes(input, 'utf-8')
            print(calculate(input, bigendian))
        else:
            # hashing mode
            input = new(mode[1], bytes(input, 'utf-8')).digest(len(input)) if 'shake' in mode[1] else new(mode[1], bytes(input, 'utf-8')).digest()
            print(calculate(input, bigendian))
    else:
        # difficulty checking op
        if len(mode) == 1:
            # default mode, e.g. no hashing
            input = bytearray.fromhex(input) if mode[0] == 'hex' else bytes(input, 'utf-8')
            print('pass' if check(int(op[1]), input, bigendian) else 'fail')
        else:
            # hashing mode
            input = new(mode[1], bytes(input, 'utf-8')).digest(len(input)) if 'shake' in mode[1] else new(mode[1], bytes(input, 'utf-8')).digest()
            print('pass' if check(int(op[1]), input, bigendian) else 'fail')


# cli
if __name__ == '__main__':
    if len(argv) < 4:
        print('calculate difficulty of a given hash or check if it reaches a threshold')
        print('usage: ' + argv[0] + ' [check:{number}|calc] [hash:{hashlib_algo}|hex|str] [utf-8 string] [bigendian=True]')
        exit(1)
    main(argv[1:])