#!/opt/local/bin/python3.6
#coding: utf-8
 
 
 
###################################################################################################
#
# Libs
#
##
 
from getopt import getopt
from base64 import b64encode as b64
from os import popen
 
from sys import argv, stdin, stdout, exit
try:
  import sha3
except:
  print("Error: 'sha3' module not found.\nYou should : pip3 install pysha3\n")
  exit(1)
 
APP=argv[0].split("/")[-1]
 
 
 
###################################################################################################
#
# Usage
#
##
 
def usage(err):
  msg ="""
Usage:
 
  %s [-a 224|256*|384|512] [-f hex|raw|b64] [file]
 
Exemples:
 
  # Compute hash of file '/bin/ls'
  %s /bin/ls
 
  # Compute hash of stdin with the content of '/bin/ls'
  cat /bin/ls|%s
 
  # Compute hash of stdin (can use `heredoc`)
  %s
 
  # Compute hash of string, with sha3_512 and display digest as base64
  %s -s 'myPasswordisverylongandsecret' -a 512 -f b64
 
Note:
 
  # SHA3 default algo is 256 (bits)
 
  # To get raw data from digest:
  printf "$(cat /bin/ls|sha3sum -a 512 -f raw)"|xxd
"""
 
  print(msg % (APP,APP,APP,APP,APP) )
  exit(err)
 
#--------------------------------------------------------------------------------------------------
 
def file_exists(filename):
  try:
    open(filename,"rb").close()
    return(True)
  except:
    return(False)
 
#--------------------------------------------------------------------------------------------------
 
def read_stdin():
  ret = stdin.buffer.read()
  return(ret)
 
#--------------------------------------------------------------------------------------------------
 
def error(txt,errcode=0):
  msg = "%s: error: "+txt
  print(msg % (APP))
  if( errcode ):
    exit(errcode)
 
#--------------------------------------------------------------------------------------------------
 
def compute_sha3(input_data,hash_algo,out_format="hex"):
  try:
    input_data = input_data.encode('utf-8')
  except:
    pass
 
  hash_algo = int(hash_algo)
 
  if( hash_algo == 224 ):
    ret = sha3.sha3_224(input_data)
  if( hash_algo == 256 ):
    ret = sha3.sha3_256(input_data)
  if( hash_algo == 384 ):
    ret = sha3.sha3_384(input_data)
  if( hash_algo == 512 ):
    ret = sha3.sha3_512(input_data)
 
  if( out_format == "hex" ):
    ret = ret.hexdigest()
  if( out_format == "b64" ):
    ret = b64( ret.digest() )
    ret = str(ret).split("'")[1]
  if( out_format == "raw" ):
    ret = ret.digest()
    ret = str(ret).split("'")[1]
 
  return( ret )
 
 
 
###################################################################################################
#
# Argments parsing
#
##
 
hash_algo  = "256"
out_format = "hex"
in_type    = "files"
item       = ""
arg_index  = 1
 
#--------------------------------------------------------------------------------------------------
 
args = argv[1:]
 
if( len(argv) >= 2 ):
  if( argv[1] == "-h" ):
    usage(0)
 
#--------------------------------------------------------------------------------------------------
 
o,a = "",""
if( len(argv) >= 2 ):
  for a in argv[1:]:
    if( o == "-a" ):
      hash_algo = a
      arg_index += 2 # -a 512
    if( o == "-f" ):
      out_format = a
      arg_index += 2 # -f b64
    o = a
 
items = argv[arg_index:]
if( len(items) == 0 ):
  in_type = "stdin"
 
 
 
 
###################################################################################################
#
# Operations
#
##
 
#--- Checks ---------------------------------------------------------------------------------------
 
if( hash_algo != "224" and hash_algo != "256" and hash_algo != "384" and hash_algo != "512" ):
  error("bad algorithm",2)
 
 
#--- Format ---------------------------------------------------------------------------------------
 
if( out_format != "hex" and out_format != "b64" and out_format != "raw" ):
  error("bad format",2)
 
 
#--- Input data -----------------------------------------------------------------------------------
 
buffers = []
src_in = "-"
 
if( in_type == "files" ):
  for item in items:
    content = ""
    errfile = False # OK
    if( file_exists(item) ):
      content = open(item,"rb").read()
    else:
      errfile = True
    buffers.append( { "content":content , "input":item , "error":errfile } )
 
elif( in_type == "stdin" ):
  buffers.append( { "content":read_stdin() , "input":"-", "error":False } )
 
else:
  error("bad input type",1)
 
 
#--- Output ---------------------------------------------------------------------------------------
 
for buff in buffers:
  if( buff["error"] == False ):
    h = compute_sha3(buff["content"],hash_algo,out_format)
    print(h + "  " + buff["input"])
 
for buff in buffers:
  if( buff["error"] == True ):
    error("unable to read file '"+buff["input"]+"'")