Brute force XOR encrypted executables using hamming distance

pe_ham_brute.py

"""
    Author: 
        Alexander Hanel
    Name: 
        pe_ham_brute.py
    Purpose:
         - POC that searches for n-grams and uses them as the XOR key.
         - Also uses hamming distance to guess key size. Check out cryptopals Challenge 6
         for more details https://cryptopals.com/sets/1/challenges/6
    Example: 
    
pe_ham_brute.py ba5aa03d724d17312d9b65a420f91285caff711e2f891b3699093cc990fdaae0
Hamming distances & calculated key sizes
[(2.6437784522003036, 58), (2.6952976867652634, 29), (3.2587556654305727, 63), (3.270363951473137, 53), (3.285315243415802, 61), (3.2863494886616276, 34), (3.29136690647482, 55), (3.300850228907783, 50), (3.306188371302278, 26), (3.309218485361723, 37)]
Length: 58, Key: IUN0mhqDx239nW3vpeL9YWBPtHC0HIUN0mhqDx239nW3vpeL9YWBPtHC0H File Name: dc53de4f4f022e687908727570345aba.bin
"""

import base64
import string
import sys
import collections
import pefile
import re
import hashlib

from cStringIO import StringIO
from collections import Counter
from itertools import cycle 
from itertools import product

DEBUG = True

def xor_mb(message, key):
    return''.join(chr(ord(m_byte)^ord(k_byte)) for m_byte,k_byte in zip(message, cycle(key)))


def hamming_distance(bytes_a, bytes_b):
    return sum(bin(i ^ j).count("1") for i, j in zip(bytearray(bytes_a), bytearray(bytes_b)))


def key_len(message, key_size):
    """"returns [(dist, key_size),(dist, key_size)]"""
    avg = []
    for k in xrange(2,key_size): 
        hd = []
        for n in xrange(len(message)/k-1):
            hd.append(hamming_distance(message[k*n:k*(n+1)],message[k*(n+1):k*(n*2)])/k)
        if hd:
            avg.append((sum(hd) / float(len(hd)), k))
    return sorted(avg)[:10]


def pe_carv(data):
    '''carve out executable using pefile's trim'''
    c = 1
    for offset in [temp.start() for temp in re.finditer('\x4d\x5a',data)]:
        # slice out executable 
        temp_buff = data[offset:]
        try:
            pe = pefile.PE(data=temp_buff)
        except:
            continue
        return pe.trim()
    return None

def write_file(data, key):
    m = hashlib.md5()
    m.update(data)
    name = m.hexdigest()
    key_name = "key-" + name + ".bin"
    file_name = name + ".bin"
    print "Length: %s, Key: %s File Name: %s" % (len(key),key, file_name)
    f =  open(file_name, "wb")
    fk = open(key_name , "wb")
    f.write(data)
    fk.write(key)
    f.close()
    fk.close()

def run(message):
    key_sizes = key_len(message, 64)
    if DEBUG:
        print "Hamming distances & calculated key sizes"
        print key_sizes
    for temp_sz in key_sizes:
        size = temp_sz[1]
        substr_counter = Counter(message[i: i+size] for i in range(len(message) - size))
        sub_count = substr_counter.most_common(32)
        for temp in sub_count:
            key, count = temp
            if count == 1:
                break
            temp = xor_mb(message, key)
            pe_c = pe_carv(temp)
            if pe_c:
                write_file(pe_c, key)
                return
    
data = open(sys.argv[1],'rb').read()
run(data)