RC4 bruteforce using entropy and file type signatures

rc4brute.py

#!/usr/bin/env python
# coding: utf-8

import sys
from Crypto.Cipher import ARC4
import numpy
import magic
import string

fname = sys.argv[1] if len(sys.argv) > 1 else 'videostream'
L_MAX = int(sys.argv[2]) if len(sys.argv) > 2 else 4

data = open(fname, 'rb').read()

#ALPH = string.digits
#ALPH = string.ascii_lowercase
#ALPH = string.ascii_uppercase
ALPH = string.letters + string.digits
#ALPH = string.letters + string.digits + string.punctuation
#ALPH = string.printable 


def combinations_with_replacement(iterable, r):
    pool = tuple(iterable)
    n = len(pool)
    if not n and r:
        return
    indices = [0] * r
    yield tuple(pool[i] for i in indices)
    while True:
        for i in reversed(range(r)):
            if indices[i] != n - 1:
                break
        else:
            return
        indices[i:] = [indices[i] + 1] * (r - i)
        yield tuple(pool[i] for i in indices)


def gen():
    for j in ALPH[len(ALPH)/2:]:
        yield tuple([j])
    for l in range(1, 3):
        for j in ALPH[len(ALPH)/2:]:
            for i in combinations_with_replacement(ALPH, l):
                yield tuple([j]) + i


def entropy(data):
    buff = numpy.frombuffer(data, dtype=numpy.uint8)
    counts = numpy.bincount(buff)
    ps = counts / float(numpy.sum(counts)) # coerce to float and normalize
    ps = ps[numpy.nonzero(ps)] # toss out zeros
    return -sum(ps * numpy.log2(ps)) # compute entropy


def check(key):
    decr = ARC4.new(''.join(key)).decrypt(data[:2**16])
    en = entropy(decr)
    sig = magic.Magic().from_buffer(decr)
    ret = magic.Magic().from_buffer(decr) != b'data' and en < 7.9
    if ret:
        sys.stdout.write('\n\n' + '='*40 + '\n' +
                         'Key: %s\n' % ''.join(key) +
                         'Sig: %s\n' % sig +
                         'Entropy: %s\n' % en +
                         '='*40 + '\n\n')
        sys.stdout.flush()
    return ret


def worker(base):
    sys.stdout.write('%s ' % ''.join(base))
    sys.stdout.flush()
    for i in combinations_with_replacement(ALPH, L_MAX):
        check(base + i)

if __name__ == "__main__":
    from multiprocessing import Pool
    p = Pool(24)
    p.map(worker, gen(), chunksize=1)
    p.close()
    p.join()