#!/usr/bin/env python
# ---
# Description: Entropy-based Analysis of PE-Sections
# Author: Philipp Winter
# Contact: pwr@7c0.org (0x532BFF15)
# License: GPLv3
#
# $Id: pytropy.py 12 2009-09-10 19:08:16Z philipp $

import sys, math, commands, os
from sys import stderr
import pefile

GNUPLOTFILE = '/tmp/gnuplot.tmp'
GNUPLOTBIN = '/usr/bin/gnuplot'

def usage( argv ):
	print >>stderr, "\nUsage: %s <file>\n" % argv

def openPortion( file=None, start=0, size=0 ):
	if file == None or size == 0:
		return ''

	fd = open(file, 'rb')
	fd.seek(start)
	data = fd.read(size)
	fd.close()

	return data

def calcEntropy( data=None ):
	if not data:
		return 0

	entropy = 0
	for char in range(256):
		p = float(data.count(chr(char))) / len(data)
		if p > 0:
			entropy += p * math.log(p, 2)
	entropy = -entropy

	return entropy

def writeDataFile( section=None, entropy=[], offset=0, size=0 ):
	fd = open("/tmp/%s.tmp" % section, 'w')
	bs = 0
	for value in entropy:
		fd.write("%s %s %s\n" % (value, str(offset + (size/2)), str(size)))
		fd.write("%s %s\n" % (offset + bs, value))
		bs = bs + 256
	fd.close()

def writeGnuplotFile( sections ):
	fd = open(GNUPLOTFILE, 'w')
	s = """
set title "Entropy-based Analysis of PE-Sections"
set grid
set style fill solid
set key box
set xrange [0:*]
set yrange [0:8]
set xlabel "File Offset (Bytes)"
set ylabel "Entropy (0.0 - 8.0)"
"""
	s += "plot '/tmp/%s.tmp' using 2:1:3 with boxes title \"%s\"" % (sections[0], sections[0])
	tmp = sections.pop(0)
	for section in sections:
		s += ", \\\n\t'/tmp/%s.tmp' using 2:1:3 with boxes title \"%s\"" % (section, section)
	sections.append(tmp)
	s += "\n"
	fd.write(s)
	fd.close()

def main( filename ):

	if not os.access(GNUPLOTBIN, os.X_OK):
		print >>stderr, "error: " + GNUPLOTBIN + " not available"
		return 1

	pe = pefile.PE(filename, fast_load=True)
	names = []
	for section in pe.sections:
		section.Name = section.Name.replace("\0", "")
		names.append(section.Name)
		entropy = [ calcEntropy(openPortion(filename, section.PointerToRawData, section.SizeOfRawData)) ]
		print "overall entropy of %s: %s" % (section.Name, entropy[0])
		print "\traw data offset: %s" % str(section.PointerToRawData)
		print "\traw data size: %s" % str(section.SizeOfRawData)
		writeDataFile(section.Name, entropy, section.PointerToRawData, section.SizeOfRawData)

	writeGnuplotFile(names)
	os.popen(GNUPLOTBIN + " -persist %s >/dev/null 2>&1" % GNUPLOTFILE)

	try:
		os.remove(GNUPLOTFILE)
		for name in names:
			os.remove("/tmp/%s.tmp" % name)
	except OSError, err:
		print err

	return 0

if __name__ == '__main__':
	if len(sys.argv) < 2 or len(sys.argv) > 3:
		usage(sys.argv[0])
		sys.exit(1)
	else:
		sys.exit(main(sys.argv[1]));