#!/usr/bin/env python

import sys,os,csv
from collections import Counter,OrderedDict
import matplotlib.pyplot as plt

if __name__ == '__main__':
  if len(sys.argv) < 2:
    sys.exit('Usage: %s malware.csv' % sys.argv[0])
  if not os.path.exists(sys.argv[1]):
    sys.exit('ERROR: CSV file %s was not found!' % sys.argv[1])

  classes = []

  # Read malware classes from a csv file (malware_hash,malware_class)
  # Create a malware count dictionary (class: number_of_samples)
  with open(sys.argv[1], 'r') as csvfile:
    content = csv.reader(csvfile, delimiter=',')
    for row in content:
      classes.append(row[1])
  counts = Counter(classes)
  print counts

  # Plot malware distribution bar chart unordered 
  plt.bar(range(len(counts)), counts.values(), align='center')
  plt.xticks(range(len(counts)), counts.keys())
  plt.show()

  # Plot malware distribution pie chart unordered
  plt.axis("equal")
  plt.pie(counts.values(),labels=counts.keys(),autopct="%1.1f%%")
  plt.title("Malware classes (%d samples)" %sum(counts.itervalues()))
  plt.show()

  # Sort malware count dictionary by number of samples
  ordcounts = OrderedDict(sorted(counts.items(), key=lambda t: t[1]))

  # Plot malware distribution bar chart ordered by number of samples 
  plt.bar(range(len(ordcounts)), ordcounts.values(), align='center')
  plt.xticks(range(len(ordcounts)), ordcounts.keys())
  plt.show()

  # Plot malware distribution pie chart ordered by number of samples
  plt.axis("equal")
  plt.pie(ordcounts.values(),labels=ordcounts.keys(),autopct="%1.1f%%")
  plt.title("Malware classes (%d samples)" %sum(counts.itervalues()))
  plt.show()
