plotframes.py

#!/usr/bin/env python

from itertools import cycle
import matplotlib.pyplot as plt
import squarify
import gzip
import sys

colormap = {}
allinsns = eval(open(sys.argv[1]).read())
mnemonics = sorted(allinsns.keys(), key=lambda k: allinsns[k])

ranges = [(int(l.split()[0]), int(l.split()[1])) for l in open(sys.argv[3])]
if ranges[-1][1] == -1:
    ranges[-1][1] = sys.maxsize

# Generate colormap
cmap = plt.get_cmap('tab20c')
colormap = dict(zip(mnemonics,cycle(cmap.colors)))

i = 0
for l in gzip.GzipFile(sys.argv[2]).readlines():
    x = l.strip().split(None,1)
    instr = int(x[0])
    
    # Limit to just the active malware ranges
    if instr < ranges[0][0]:
        continue
    elif ranges[0][0] <= instr <= ranges[0][1]:
        pass
    else:
        try:
            ranges.pop(0)
        except IndexError:
            # Done with the last range, we can quit
            break
        # Check if current record falls inside the new range
        if instr < ranges[0][0]:
            continue
    print("Current range: %s" % str(ranges[0]))

    d = eval(x[1])
    labels, values = list(zip(*d.items()))
    colors = list(map(lambda l: colormap[l],labels))
    squarify.plot(sizes=values,label=labels,color=colors)
    plt.axis('off')
    plt.savefig('%08d.png' % i)
    print(instr, '%08d.png' % i)
    plt.clf()
    plt.close()
    i += 1