From afc4a84e57eba1f4a736076aa0a1856b095df69a Mon Sep 17 00:00:00 2001
From: dejanlevaja <dejan@levaja.com>
Date: Tue, 15 Mar 2016 15:16:06 +0100
Subject: [PATCH] v3.7 MAC loookup optimization

---
 main.py          | 73 ++++++------------------------------------------
 util_uploader.py | 15 ++++++++++
 2 files changed, 23 insertions(+), 65 deletions(-)

diff --git a/main.py b/main.py
index 8f85929..d4567b8 100644
--- a/main.py
+++ b/main.py
@@ -16,56 +16,13 @@
 import module_openbsd as openbsd
 import module_aix as aix
 
-__version__ = "3.6"
+__version__ = "3.7"
 
 # environment and other stuff
 lock = threading.Lock()
 q = Queue.Queue()
 
-
-def mac_lookup_map():
-    rest = uploader.Rest(base_url, username, secret, debug)
-    url = base_url + '/api/1.0/macs/'
-    response = rest.fetcher(url)
-
-    #print json.dumps(r, indent=4, sort_keys=True)
-    if  isinstance(response, dict):
-        macs = response['macaddresses']
-        mac_devid_map = {}
-
-        # single device might have multiple macs
-        # create dict to store dev_id:{mac_id:mac} mappings
-        macdb = {}
-        for m in macs:
-            try:
-                mid = int(m['macaddress_id'])
-                mac = m['macaddress']
-                dev_id = m['device']['device_id']
-                if dev_id not in macdb:
-                    macdb.update({dev_id:[{mid:mac}]})
-                else:
-                    macdb[dev_id].append({mid:mac})
-            except:
-                pass
-        if macdb:
-            for k,values in macdb.items():
-                # we must discover smallest mac_id. Mac with smallest mac_id is the 'first' one and it is used for mac lookup
-                smallest = 99999999999999999999
-                mac = None
-                dev_id = k
-                for value in values:
-                    mid = value.keys()[0]
-                    if mid < smallest:
-                        smallest = mid
-                        mac = value[mid]
-                # create mac:dev_id mapping
-                mac_devid_map.update({mac:dev_id})
-
-            macdb.clear()
-        return mac_devid_map
-
-
-def find_devid_by_mac(data):
+def find_devid_by_mac(data, rest):
     macs = []
     for rec in data:
         if 'macaddress' in rec:
@@ -73,22 +30,17 @@ def find_devid_by_mac(data):
             macs.append(m)
 
     for mac in macs:
-        if mac in mac_devid_map:
-            dev_id = mac_devid_map[mac]
+        dev_id = rest.get_device_by_mac(mac)
+        if dev_id:
             return dev_id
 
-
-
 def upload(data):
-    dev_id = None
-    if mac_lookup:
-        dev_id = find_devid_by_mac(data)
-
-
     ips = []
     name = None
-    result = None
+    dev_id = None
     rest = uploader.Rest(base_url, username, secret, debug)
+    if mac_lookup:
+        dev_id = find_devid_by_mac(data, rest)
 
     # get hdd parts if any
     hdd_parts = {}
@@ -415,8 +367,6 @@ def check_os(ip):
 
 
 def main():
-    if mac_lookup:
-        mac_lookup_map()
     sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
     sock.settimeout(float(timeout))
     msg = '\r\n[!] Running %s threads.' % THREADS
@@ -469,16 +419,9 @@ def main():
 
 if __name__ == '__main__':
     from module_shared import *
-    if mac_lookup:
-        if debug:
-            print '\n[!] Creating MAC<->deviceID map for MAC lookup.\n'
-        mac_devid_map = mac_lookup_map()
     main()
     sys.exit()
 else:
     # you can use dict_output if called from external script (starter.py)
     from module_shared import *
-    if mac_lookup:
-        if debug:
-            print '\n[!] Creating MAC<->deviceID map for MAC lookup.\n'
-        mac_devid_map = mac_lookup_map()
+
diff --git a/util_uploader.py b/util_uploader.py
index 01cf888..ebf46be 100644
--- a/util_uploader.py
+++ b/util_uploader.py
@@ -136,3 +136,18 @@ def delete_ip(self, ip):
             for ip_id in ip_ids:
                 url = self.base_url + '/api/1.0/ips/%s' % ip_id
                 self.deleter(url)
+
+    def get_device_by_mac(self, mac):
+        if not DRY_RUN:
+            url = self.base_url + '/api/1.0/macs/?mac=%s' % mac
+            msg = '\r\nFind device by mac:  %s ' % mac
+            if self.debug:
+                print msg
+            response = self.fetcher(url)
+            if isinstance(response, dict) and 'macaddresses' in response:
+                dev_id = [x['device']['device_id'] for x in response['macaddresses'] if 'device' in x]
+                if dev_id:
+                    try:
+                        return dev_id[0]
+                    except:
+                        pass
\ No newline at end of file