#!/usr/bin/python

# Demo Script - Testing Matplotlib and lxml
# Author: Tyler Reguly - ht [at] computerdefense [dot] org

from lxml import etree
from pylab import figure, bar, yticks, xticks, FixedFormatter, gca, savefig, ylabel, xlabel, title
from numpy import arange
from sys import argv, exit

def usage():
    print "Demo Script to Determine the Frequency of Ports from a nmap XML file"
    print "Created by Tyler Reguly - ht [at] computerdefense [dot] org"
    print
    print "usage: %s <infile> <outfile>" % argv[0]
    print "infile: nmap XML"
    print "outfile: <file>.png"
    exit()


ports = {}
values = []
keys = []

try: fileIn = argv[1]
except IndexError:
    usage()
try: fileOut = argv[2]
except IndexError:
    usage()

fh = open(fileIn)

for event, element in etree.iterparse(fh):
    if element.tag == 'port':
        if ports.has_key(int(element.attrib['portid'])):
            ports[int(element.attrib['portid'])] +=1 
        else:
            ports[int(element.attrib['portid'])] = 1
portList = ports.keys()
portList.sort()
for key in portList: 
    keys.append(key)
    values.append(ports[key])
index = arange(len(values))
width = 0.5
bar(index, values, width, color = 'red')
tick_range = arange(max(values)+2)
yticks(tick_range, size='xx-small')
xticks(index+width, keys, size='xx-small', rotation='vertical')
yformatter = FixedFormatter([str(x) for x in tick_range])
xformatter = FixedFormatter([str(x) for x in keys])
gca().yaxis.set_major_formatter(yformatter)
gca().xaxis.set_major_formatter(xformatter)
ylabel('Frequency')
xlabel('Ports')
title('Frequency of Ports in nmap XML')
savefig(fileOut)