#!/usr/bin/env python
import sys
import os
import getopt
from time import time

sys.path.append(os.path.join(sys.path[0],'lib'))

'''
__author__="David Busby"
__copyright__="David Busby <oneiroi@fedorapeople.org>"
__license__="GNU v3 + part 5d section 7: Redistribution/Reuse of this code is permitted under the GNU v3 license, as an additional term ALL code must carry the original Author(s) credit in comment form."
'''

'''
Nyancat eatz yummy passwords and poops out rainbows!
'''

try:
    import multiprocessing
except:
    print 'This script requires the multiprocessing lib to function, included in python 2.6+ or available from pypi for 2.5'
    sys.exit(1)

class opts:
    slen = 0

def usage():
    print sys.argv[0],'-f /path/to/wordlist.txt -o /output/path/to/rainbow_table.txt -r hash_method -t <max threads, default 1>'

def progress(str):
        #@todo: rainbow asii animated progress nyan!
        str = " %s" % str

        while len(str) < opts.slen:
            str = '%s ' % str
        opts.slen = len(str)
        sys.stdout.write(str + '\r')
        sys.stdout.flush()

def main():
    try:
        opts,args = getopt.getopt(sys.argv[1:],'f:o:r:t:',[])
    except getopt.GetoptError, e:
        print e
        usage()
        sys.exit(2)
    
    wlist   = None
    ofile   = None
    hMethod = None
    threads = 2

    for o,a in opts:
        if o == '-f':
            wlist = a
        elif o == '-o':
            ofile = a
        elif o == '-t':
            threads = int(a)
        elif o == '-r':
            hMethod = a
            try:
                exec('import lib.%s as rainbow'%a)
            except ImportError as e:
                print e
                sys.exit(1)   
        else:
            assert False,'Unsupported option %s' %a
    
    if ofile == None:
        print 'Output file not specified, exiting'
        sys.exit(2)
    if wlist == None:
        print 'Wordlist file not specified exiting'
        sys.exit(2)

    p = multiprocessing.Pool(processes=threads) 
    i = 0
    cur = 0
    words = []
    print """
+        o     o       +        o
-_-_-_-_-_-_-_,------,      o 
_-_-_-_-_-_-_-|   /\_/\  
-_-_-_-_-_-_-~|__( ^ .^)  +     +  
_-_-_-_-_-_-_-""  ""      
"""
    print 'Now streaming words from %s to rainbow file %s' %(wlist,ofile)
    ofp = open(ofile, 'w+')
    wsize = os.path.getsize(wlist)
    wfp = open(wlist,'r')
    sTime = time()
    for word in wfp:
        words.append(word)
        i+=1
        cur+=1
        #progress('%d words'%i)
        #if we have reached a decent data set chunk size, lets process it now and stream out to file
        #@todo: it will take some expeimentation to tune this value, too low and exesive context switching due to thread spawn / death, too high and the memory usage may lead to swapping.
        #@todo: Add memory option, i.e. -m 512 == 512mb is the upper limit at which point threads are spawned.
        if cur >= threads*100000:
            progress('Rainbow table generation in progress: %.f%%' % ((wfp.tell()*1.00/wsize*1.00)*100))
            data=p.map(rainbow.hash,words)
            for item in data:
                for h in item:
                    ofp.write('%s:%s\n' % (h,item[h]))
            del words
            del cur
            cur = 0
            words = []
    if words.__len__() > 0:
      data=p.map(rainbow.hash,words)
      for item in data:
        for h in item:
          ofp.write('%s:%s\n' % (h,item[h]))
    ofp.close()
    eTime = time()
    print
    print 'Completed %s hashing of %d words in %.2f seconds (%.2f/s) using %d threads' % (hMethod,i,(eTime-sTime),(i/(eTime-sTime)),threads)

if __name__ == '__main__':
    main()
