#!/usr/bin/env python

import sys
import string
from work_queue import *

# CHARSET = string.printable
CHARSET = string.ascii_lowercase
WRAPPER = '''
#!/bin/bash

host=$(hostname --fqdn)
if [ -z "${host/%*.cluster/}" ]; then
  module load python/2.7 > /dev/null 2> /dev/null
fi
"$@"
'''

sum = sys.argv[1]

try:
    Q = WorkQueue(port = 0, name = "fg-tutorial", catalog = True)
except:
    print "could not instantiate Work Queue master"
    sys.exit(1)

print "Listening on port %d." % Q.port

def strproduct (charset, repeat):
    def helper (charset, repeat, list):
        if repeat == 0:
            yield string.join(list, "")
        else:
            for char in charset:
                list.append(char)
                for product in helper(charset, repeat-1, list):
                    yield(product)
                list.pop()
    list = []
    return helper(charset, repeat, list)

def task_generator (Q, sum):
    for i in range(0, 32):
        for prefix in strproduct(CHARSET, i):
            prefix = string.join(prefix, '')
            print("Dispatching task with prefix '%s'." % prefix)
            T = Task("/bin/bash wrapper.sh python md5-brute.py '%s' '%s'" % (prefix, sum))
            T.specify_buffer(WRAPPER, "wrapper.sh", cache = True)
            T.specify_input_file("md5-brute.py", cache = True)
            taskid = Q.submit(T)
            yield taskid

for taskid in task_generator(Q, sum):
    while not Q.hungry():
        T = Q.wait(5)
        if T:
            if T.return_status != 0:
                print(T.output)
                sys.exit(T.return_status)
            elif len(T.output) > 0:
                print("Password is '%s'." % T.output)
                sys.exit(0)