1#!/usr/bin/env python
2__author__ = 'Ben "TheX1le" Smith, Marfi'
3__email__ = 'thex1le@gmail.com'
4__website__= ''
5__date__ = '04/26/2011'
6__version__ = '2011.4.26'
7__file__ = 'ouiParse.py'
8__data__ = 'a class for dealing with the oui txt file'
9
10"""
11########################################
12#
13# This program and its support programs are free software; you can redistribute it and/or modify it
14# under the terms of the GNU General Public License version 2 as
15# published by the Free Software Foundation; version 2.
16#
17# This program is distributed in the hope that it will be useful, but
18# WITHOUT ANY WARRANTY; without even the implied warranty of
19# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
20# General Public License for more details.
21#
22#########################################
23"""
24
25import re
26import urllib
27import sys
28import os
29import pdb
30#this lib is crap and needs to be rewritten -Textile
31
32if os.path.isdir('./support/'):
33    path='./support/'
34elif os.path.isdir('/usr/local/share/airgraph-ng/'):
35    path='/usr/local/share/airgraph-ng/'
36elif os.path.isdir('/usr/share/airgraph-ng/'):
37    path='/usr/share/airgraph-ng/'
38else:
39    raise Exception("Could not determine path, please, check your installation")
40
41class macOUI_lookup:
42    """
43    A class for deaing with OUIs and deterimining device type
44    """
45    def __init__(self, oui=False):
46        """
47        generate the two dictionaries and return them
48        """
49        #a poor fix where if we have no file it trys to download it
50        self.ouiTxtUrl   = "http://standards-oui.ieee.org/oui.txt"
51
52        self.ouiTxt = oui
53        if not oui or not os.path.isfile(self.ouiTxt):
54            self.ouiUpdate()
55            self.ouiTxt = path + "oui.txt"
56        self.last_error = None
57        self.identDeviceDict(path + 'ouiDevice.txt')
58        self.identDeviceDictWhacMac(path + 'whatcDB.csv')
59        self.ouiRaw      = self.ouiOpen(self.ouiTxt)
60        self.oui_company = self.ouiParse()  #dict where oui's are the keys to company names
61        self.company_oui = self.companyParse()  #dict where company name is the key to oui's
62
63    def compKeyChk(self,name):
64        """
65        check for valid company name key
66        """
67        compMatch = re.compile(name,re.I)
68        if self.company_oui.has_key(name):
69            return True
70        for key in self.company_oui.keys():
71                if compMatch.search(key) is not None:
72                    return True
73        return False
74
75    def ouiKeyChk(self,name):
76        """
77        check for a valid oui prefix
78        """
79
80        if self.oui_company.has_key(name):
81            return True
82        else:
83            return False
84
85    def lookup_OUI(self,mac):
86        """
87        Lookup a oui and return the company name
88        """
89        if self.ouiKeyChk(mac) is not False:
90            return self.oui_company[mac]
91        else:
92            return False
93
94    def lookup_company(self,companyLst):
95        """
96        look up a company name and return their OUI's
97        """
98        oui = []
99        if type(companyLst) is list:
100            for name in companyLst:
101                compMatch = re.compile(name,re.I)
102                if self.company_oui.has_key(name):
103                    oui.extend(self.company_oui[name])
104                else:
105                    for key in self.company_oui:
106                        if compMatch.search(key) is not None:
107                            oui.extend(self.company_oui[key])
108
109        elif type(companyLst) is str:
110            if self.company_oui.has_key(companyLst):
111                oui = self.company_oui[companyLst]
112            else:
113
114                compMatch = re.compile(companyLst,re.I)
115                for key in self.company_oui:
116                    if compMatch.search(key) is not None:
117                        oui.extend(self.company_oui[key]) #return the oui for that key
118        return oui
119
120    def ouiOpen(self,fname,flag='R'):
121        """
122        open the file and read it in
123        flag denotes use of read or readlines
124        """
125        try:
126            with open(fname, "r") as fid:
127                if flag == 'RL':
128                    text = fid.readlines()
129                elif flag == 'R':
130                    text = fid.read()
131            return text
132        except IOError:
133            return False
134
135
136    def ouiParse(self):
137        """
138        generate a oui to company lookup dict
139        """
140        HexOui= {}
141        Hex = re.compile('.*(hex).*')
142        #matches the following example "00-00-00   (hex)\t\tXEROX CORPORATION"
143        ouiLines = self.ouiRaw.split("\n")
144        #split each company into a list one company per position
145        for line in ouiLines:
146            if Hex.search(line) is not None:
147                lineList = Hex.search(line).group().replace("\t"," ").split("  ")
148                #return the matched text and build a list out of it
149                HexOui[lineList[0].replace("-",":")] = lineList[2].strip()
150                #build a dict in the format of mac:company name
151        return HexOui
152
153    def companyParse(self):
154        """
155        generate a company to oui lookup dict
156        """
157        company_oui = {}
158        for oui in self.oui_company:
159            if company_oui.has_key(self.oui_company[oui]):
160                company_oui[self.oui_company[oui]].append(oui)
161            else:
162                company_oui[self.oui_company[oui]] = [oui]
163        return company_oui
164
165
166    def ouiUpdate(self):
167        """
168        Grab the oui txt file off the ieee.org website
169        """
170        try:
171            print("Getting OUI file from %s to %s" %(self.ouiTxtUrl, path))
172            urllib.urlretrieve(self.ouiTxtUrl, path + "oui.txt")
173            print "Completed Successfully"
174        except Exception, error:
175            print("Could not download file:\n %s\n Exiting airgraph-ng" %(error))
176            sys.exit(0)
177
178    def identDeviceDict(self,fname):
179        """
180        Create two dicts allowing device type lookup
181        one for oui to device and one from device to OUI group
182        """
183        self.ouitodevice = {}
184        self.devicetooui = {}
185        data = self.ouiOpen(fname,'RL')
186        if data == False:
187            self.last_error = "Unable to open lookup file for parsing"
188            return False
189        for line in data:
190            dat = line.strip().split(',')
191            self.ouitodevice[dat[1]] = dat[0]
192            if dat[0] in self.devicetooui.keys():
193                self.devicetooui[dat[0]].append(dat[1])
194            else:
195                self.devicetooui[dat[0]] = [dat[1]]
196
197    def identDeviceDictWhacMac(self,fname):
198        """
199        Create two dicts allowing device type lookup from whatmac DB
200        one for oui to device and one from the device to OUI group
201        """
202        self.ouitodeviceWhatmac3 = {}
203        self.ouitodeviceWhatmac = {}
204        self.devicetoouiWhacmac = {}
205        data = self.ouiOpen(fname,'RL')
206        if data == False:
207            self.last_error = "Unble to open lookup file for parsing"
208            return False
209        for line in data:
210            dat = line.strip().split(',')
211            dat[0] = dat[0].upper()
212            self.ouitodeviceWhatmac[dat[0]] = dat[1]
213            self.ouitodeviceWhatmac3[dat[0][0:8]] = dat[1] # a db to support the 3byte lookup from whatmac
214            if dat[1] in self.devicetoouiWhacmac.keys():
215                self.devicetoouiWhacmac[dat[1]].append(dat[0])
216            else:
217                self.devicetoouiWhacmac[dat[1]] = [dat[0]]
218
219