# python3 script to find PIN corresponding to a hash, using multiple processes on one node
#
# R. Perry, August 2018
#
from hashlib import sha1
from time import time
from multiprocessing import Process, Queue, current_process, cpu_count

# check hash for n-digit PINs from start to stop-1, sending 4 progress messages along the way
#
def search( n, start, stop, h, q):
  name = current_process().name; fmt = "0" + str(n) + "d";
  M4 = (stop-start)//4;  M = start + M4;  b = "   "
  q.put( name + " " + format(start,fmt) + " ...")
  pin = start
  while pin < stop:
    if pin == M: q.put( name + b + " ... " + format(pin,fmt));  M += M4;  b += "   "
    d = bytearray( format(pin,fmt), 'utf-8')
    md = sha1(d).digest()
    if md == h: q.put( name + " match: PIN = " + format(pin,fmt) ); break
    pin += 1
  q.put( name + b + " ... " + format(pin-1,fmt) + " done") 


if __name__ == '__main__':

  # set hash, number of PIN digits, and number of processes
  #
  hash = input( "hash in hex: ");  h = bytes.fromhex( hash)

  n = int( input( "number of PIN digits: "));  N = pow( 10, n)

  print( "cpu_count =", cpu_count());  np = int( input( "number of processes: "))

  print( "Checking", n, "digit PINs,", N, "combinations,", np, "processes")
  print( "to match hash", hash)

  t0 = time() # save start time

  # set up processes and result queue
  #
  q = Queue()
  d = N//np # might not divide evenly
  p = [None for i in range(np)]
  for i in range(np):
    start = d*i;  stop = d*(i+1) if i < np-1 else N # handle d roundoff
    p[i] = Process( target=search, args=(n, start, stop, h, q,))
    p[i].start()
 
  # get results
  #
  done = 0
  while done < np:
    msg = q.get();  print(msg);
    if "match" in msg: break
    if "done" in msg: done += 1

  # if match found, terminate the processes
  #
  if done < np:
    for i in range(np): p[i].terminate()
  else:
    print( "\nNo match found.")

  t1 = time() # end time

  print( "\nElapsed time:", '%g' % (t1-t0), "seconds\n")

  # input( "Press Enter to exit:")