#!/usr/bin/env python
#python 3.0 support
from __future__ import print_function

__author__ = 'Ben "TheX1le" Smith'
__email__  = 'thex1le _A_T_ remote-exploit.org'
__website__= 'remote-exploit.org'
__date__   = '05/18/2011'
__version__= '2.0.1'
__file__   = 'airgraph-ng'
__data__   = 'This is the main airgraph-ng file'
import subprocess
import sys, optparse, os
import pdb

try:
    from graphviz import *
except ImportError, error:
    raise Exception("Your airgraph-ng installation is broken, could not import libraries: %s" %(error))

path = libOuiParse.path

class interface:
    """
    provide basic UI to user
    """
    def header(self):
        """
        Print a pretty header out
        """
        print('#'*42+"\n#"+" "*9+"Welcome to Airgraph-ng"+" "*9+"#\n"+"#"*42+"\n")

class dotCreate:
    """
    Class for creating graphviz .dot config files
    """
    def __init__(self, file_):
        ret = libDumpParse.airDumpParse().parser(file_)
        self.capr   = ret['capr']
        self.Ap     = ret['apDict']
        self.client = ret['clientDict'] 
        self.NA     = ret['NA']
        self.NAP    = ret['NAP']

        #144x144 hard code image size to 12feet x 12feet
        #start graphviz config file
        self.dotFile = ['digraph G{\n\tsize ="144,144";\n\toverlap=false;\n']
        self.ouiCheck = libOuiParse.macOUI_lookup(path + 'oui.txt')
        if self.ouiCheck is False:
            print("Missing the oui.txt file from http://standards.ieee.org/develop/regauth/oui/oui.txt, place it in the support directory")
            sys.exit(-1)
        self.ouiCheck.identDeviceDict(path + 'ouiDevice.txt')

    def CAPR(self):
        """
        Client AP relationship graph
        Display a graph showing what clients are talking to what AP's
        """
        tclient = 0 #total client number
        tap     = len(self.capr.keys()) #total ap list
        for bssid in self.capr.keys():
            time = [self.Ap[bssid]['lts'],self.Ap[bssid]['fts']]
            priv = self.Ap[bssid]['privacy']
            if priv == '':
                priv = "Unkown"
            Color = self.encColor(priv)
            tclient += len(self.capr[bssid])
            for client in self.capr[bssid]:
                self.dotFile.append(self.linker(bssid,'->',client))
                self.dotFile.append(self.clientColor(client,'black',[self.client[client]['lts'],self.client[client]['fts']],None,True))
            self.dotFile.append(
                self.apLabel(bssid,Color,len(self.capr[bssid]),time))
        footer ='label="Generated by Airgraph-ng\\n%s Access Points and\\n%s Clients shown";\n}' %(tap,tclient)
        self.dotFile.append(footer)
    
    def CPG(self):
        """
        Common Probe Graph
        Shows a graph of every client requesting similar probes
        """
        probeCount = 0
        clientCount = {}
        for key in self.client.keys():
            sdata = self.client[key]
            lts = sdata['lts']
            fts = sdata['fts'] 
            if sdata["probe"] != ['']:
                lpc = len(sdata['probe'])
                clientCount[key] = key
                probeCount += lpc
                for probe in sdata['probe']:
                    self.dotFile.append(self.clientColor(probe,'blue'))
                    self.dotFile.append(self.linker(key,'->',probe))
                    clientColor = '%s\\nRequested %s probes' %(key,lpc)
                    self.dotFile.append(self.clientColor(key,'black',[lts,fts],None,True))
        footer = 'label="Generated by Airgraph-ng\\n%s Probes and \\n%s Clients are shown";\n}' % (probeCount,len(clientCount.keys()))
        self.dotFile.append(footer)

    def clientColor(self,mac,color,time=None,label=None,DouiCheck=False):
        """
        format the client with a color and a label
        return graphiz format line
        """
        
        if label == None:
            label = mac
        #device OUI check
        if DouiCheck == True:
            lts = time[0]
            fts = time[1]
            rtn = '\t"%s" [label="%s\\nOUI: %s\\nDevice Type: %s\\nFirst Time Seen: %s\\nLast Time Seen: %s" color=%s fontsize=9];\n' %(mac,mac,self.APouiLookup(mac),self.clientOuiLookup(mac),fts,lts,color)
        else:
            rtn = '\t"%s" [label="%s" color=%s fontsize=9];\n' %(mac,label,color)

        return rtn

    def encColor(self,enc):
        """
        Take encryption type and decide what color it should be displayed as
        returns a list containing AP fill color and Ap label font color
        """
        fontColor = "black" #default font color

        if enc =="OPN":
            color = "firebrick2"
        elif enc == "WEP":
            color = "gold2"
        elif enc in ["WPA","WPA2WPA","WPA2","WPAOPN"]:
            color = "green3"
        else: #unknown enc type
            color = "black"
            fontColor = "white"

        return (color,fontColor)

    def linker(self,objA,sep,objB):
        """
        Retrun a graphviz line that links 2 objects togeather
        Both objects are passed in with a separator
        returns graphiz format line
        """
        return '\t"%s"%s"%s";\n' %(objA,sep,objB)

    def dotWrite(self):
        """
        Write all the information obtained to a config file
        """
        dotdata = ''.join(self.dotFile)
        try:
            os.remove('airGconfig.dot')
        except Exception:
            pass
        nfile = open('airGconfig.dot','a')
        nfile.write(dotdata)
        nfile.close()
    
    def clientOuiLookup(self,oui):
        """
        check ouiDevices and attempt to determine the device type
        """
        prefix = oui[:8]
        if prefix in self.ouiCheck.ouitodevice:
            device = self.ouiCheck.ouitodevice[prefix]
        else:
            device = 'Unknown'

        return device

    def APouiLookup(self,oui):
        """
        check the oui.txt file and determine the manufacture for an AP or client
        """
        prefix = oui[:8]
        company = self.ouiCheck.lookup_OUI(prefix)
        if company == False:
            company = 'Unknown'
        return company

    def apLabel(self,bssid,color,cnum,time):
        """
        Create label strings for AP's
        """
        lts = time[0]
        fts = time[1]
        AP = self.Ap[bssid]
        return'\t"%s" [label="%s\\nEssid: %s\\nChannel: %s\\nEncryption: %s\\nOUI: %s\\nFirst Time Seen: %s\\n Last Time Seen: %s\\nNumber of Clients: %s" style=filled fillcolor="%s" fontcolor="%s" fontsize=9];\n' %(bssid,bssid,AP['essid'].rstrip('\x00'),AP['channel'],AP['privacy'],self.APouiLookup(bssid),fts,lts,cnum,color[0],color[1])

    def graphCreate(self,fname,outname):
        """
        Write out the graphviz dotFile and creat the graph
        """
        print("\n**** WARNING Images can be large, up to 12 Feet by 12 Feet****")
        print("Creating your Graph using, %s and writing to, %s"  %(fname,outname))
        print("Depending on your system this can take a bit. Please standby......")
    
        try:
            subprocess.Popen(["fdp","-Tpng","airGconfig.dot","-o",outname,"-Gcharset=latin1"]).wait()
        except Exception:
            os.remove("airGconfig.dot")
            print("You seem to be missing the Graphviz toolset. Did you check out the airgraph-ng Deps in the readme?")
            sys.exit(1)
        os.remove("airGconfig.dot") #comment this line out for dotfile debugging

if __name__ == "__main__":
    """
    Main function.
    Parses command line input for proper switches and arguments. 
    Error checking is done in here.
    Variables are defined and all calls are made from MAIN.
    """
    parser = optparse.OptionParser("usage: %prog options [-o -i -g ] ")  #
    parser.add_option("-o", 
        "--output", 
        dest="output",
        nargs=1, 
        help="Our Output Image ie... Image.png")
    parser.add_option("-i", 
        "--dump", 
        dest="input", 
        nargs=1 
        ,help="Airodump txt file in CSV format NOT the pcap")
    parser.add_option("-g", 
        "--graph",
        dest="graph_type", 
        nargs=1,
        help="Graph Type Current [CAPR (Client to AP Relationship) OR CPG (Common probe graph)]")
    
    if len(sys.argv) <= 1:
        interface().header()
        parser.print_help()
        sys.exit(0)
    (options, args) = parser.parse_args()

    outFile = options.output
    graphType = options.graph_type
    inFile = options.input

    if inFile == None:
        print("Error No Input File Specified")
        sys.exit(1)
    if outFile == None: 
        outFile = options.input.replace('.txt', '.png')
    if graphType.upper() not in ['CAPR','CPG','ZKS']:
        print("Error Invalid Graph Type\nVaild types are CAPR or CPG")
        sys.exit(1)
    if graphType == None:
        print("Error No Graph Type Defined")
        sys.exit(1)
    dot = dotCreate(inFile)
    if graphType.upper() == 'CPG':
        dot.CPG()
        dot.dotWrite()
        dot.graphCreate(inFile,outFile)
    elif graphType.upper() == 'CAPR':
        dot.CAPR()
        dot.dotWrite()
        dot.graphCreate(inFile,outFile)
