Some work on data, repr() and cleaning

3 files changed, 57 insertions, 58 deletions

diff --git a/haircontrol/data.py b/haircontrol/data.py
index 599d322..d8195c7 100644
--- a/haircontrol/data.py
+++ b/haircontrol/data.py
@@ -1,3 +1,30 @@
+class EtherDomain:
+    def __init__(self):
+        self.equipments = {}
+        self.ip2mac = {}
+        self.mac2ip = {}
+
+    def __repr__(self):
+        e_list = list(self.equipments.values())
+        e_list.sort(key=lambda x: x.name)
+        return '([\n%s\n],\n%s\n)'%(',\n'.join('  %s'%repr(e) for e in e_list), repr(self.ip2mac))
+
+    def add_equipment(self,e):
+        old = self.equipments.get(e.mgmtip)
+        if old:
+            print("Warn : %s replaced by %s"%(old,e))
+        self.equipments[e.mgmtip] = e
+
+    def index_mac_ip(self,mac, ip):
+        oldmac = self.ip2mac.get(ip)
+        if oldmac:
+            print("Warn : %s replaced by %s for %s"%(oldmac, mac, ip))
+        oldip = self.mac2ip.get(mac)
+        if oldip:
+            print("Warn : %s replaced by %s for %s"%(oldip, ip, mac))
+        self.ip2mac[ip] = mac
+        self.mac2ip[mac] = ip
+
 class Equipment:
     def __init__(self, name=None, mgmtip=None):
         self.name = name
@@ -5,13 +32,12 @@ class Equipment:
         self.ifaces = {}
 
     def __repr__(self):
-        return repr( (self.name, self.mgmtip, self.ifaces.values()) )
+        return repr( (self.name, self.mgmtip, list(self.ifaces.values())) )
 
     def add_seen_mac(self, ifname, mac):
         iface = self.ifaces.get(ifname)
         if not iface:
-            iface = Interface()
-            iface.name = ifname
+            iface = Interface(ifname)
             self.ifaces[ifname] = iface
         iface.mac_seen.append(mac)
 
@@ -20,8 +46,7 @@ class Interface:
         self.name = name
         self.mac = mac
         self.mac_seen = []
-        self.remote = []
 
     def __repr__(self):
-        return repr( (self.mac, self.name) )
+        return repr( ( self.mac, self.name, '[ %i mac_seen ]'%len(self.mac_seen) ) )
 
diff --git a/haircontrol/discovery.py b/haircontrol/discovery.py
index d454283..2b7ad83 100644
--- a/haircontrol/discovery.py
+++ b/haircontrol/discovery.py
@@ -1,33 +1,20 @@
 
 import re
 import xml.etree.ElementTree
-from haircontrol.data import Equipment
+from haircontrol.data import *
 
 class Discovery:
+    IPNEIGH = re.compile("^(?P<ip>[a-f0-9:.]+) dev (?P<ifname>.*) lladdr (?P<mac>[a-f0-9:]*)")
+    # fe80::8300 dev eth1 lladdr 10:fe:ed:f1:e1:f3 router STALE
+    # 172.16.20.210 dev eth1 lladdr c0:4a:00:fe:1f:87 REACHABLE
+
     def __init__(self, inspector):
         self.inspector = inspector
-        self.equipments = {}
-        self.ip2mac = {}
-        self.mac2ip = {}
-
-    def add_equipment(self,e):
-        old = self.equipments.get(e.mgmtip)
-        if old:
-            print("Warn : %s replaced by %s"%(old,e))
-        self.equipments[e.mgmtip] = e
-
-    def index_mac_ip(self,mac, ip):
-        oldmac = self.ip2mac.get(ip)
-        if oldmac:
-            print("Warn : %s replaced by %s for %s"%(oldmac, mac, ip))
-        oldip = self.mac2ip.get(mac)
-        if oldip:
-            print("Warn : %s replaced by %s for %s"%(oldip, ip, mac))
-        self.ip2mac[ip] = mac
-        self.mac2ip[mac] = ip
+        self.net = EtherDomain()
 
-    def discover_hinting_from_lldp(self, lldpnode):
-        self.inspector.connect(lldpnode)
+    def discover_hinting_from_lldp(self, e_lldp):
+        self.net.add_equipment(e_lldp)
+        self.inspector.connect(e_lldp)
         fd = self.inspector.command('lldp')
         root = xml.etree.ElementTree.parse(fd).getroot()
         for iface in root.iter('interface'):
@@ -36,33 +23,23 @@ class Discovery:
             e = Equipment()
             e.name = chassis.find('name').text
             e.mgmtip = chassis.find('mgmt-ip').text
-            self.add_equipment(e)
+            self.net.add_equipment(e)
             for port in iface.findall('port'):
                 #remote_ifname = port.find('id').text
                 e.add_seen_mac(ifname, 'lldp') #XXX
         fd.close()
         self.inspector.disconnect()
 
-    def discover_from_root(self, rootnode):
-        self.inspector.connect(rootnode)
+    def discover_from_root(self, e_root):
+        self.net.add_equipment(e_root)
+        self.inspector.connect(e_root)
         fd = self.inspector.command('ip-neigh')
-        IPNEIGH = re.compile("^(?P<ip>[a-f0-9:.]+) dev (?P<ifname>.*) lladdr (?P<mac>[a-f0-9:]*)")
         for line in fd:
-            matches = IPNEIGH.search(line)
+            matches = Discovery.IPNEIGH.search(line)
             if matches:
                 ip, ifname, mac = [ matches.group(k) for k in ['ip','ifname','mac'] ]
-                self.index_mac_ip(mac, ip)
-                rootnode.add_seen_mac(ifname, mac)
+                self.net.index_mac_ip(mac, ip)
+                e_root.add_seen_mac(ifname, mac)
         fd.close()
         self.inspector.disconnect()
 
-# fe80::8300 dev eth1 lladdr 10:fe:ed:f1:e1:f3 router STALE
-# 172.16.11.46 dev eth1 lladdr 24:a4:3c:ee:89:ca STALE
-# 172.16.20.3 dev eth1 lladdr 00:27:22:0e:74:15 STALE
-# 172.16.20.210 dev eth1 lladdr c0:4a:00:fe:1f:87 REACHABLE
-# 172.16.21.69 dev eth1 lladdr e8:de:27:b5:f2:b1 DELAY
-# 172.16.20.216 dev eth1 lladdr c0:4a:00:fe:09:bd PERMANENT
-# 172.16.11.41 dev eth1 lladdr 04:18:d6:0e:37:d4 STALE
-# 172.16.11.104 dev eth1 lladdr 00:15:6d:8e:22:46 STALE
-# 172.16.10.8 dev eth1 lladdr 00:27:22:0e:67:f9 STALE
-# 
diff --git a/tests/test_discovery.py b/tests/test_discovery.py
index 96d262e..b0de19b 100755
--- a/tests/test_discovery.py
+++ b/tests/test_discovery.py
@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 import unittest
-import pprint
+#import pudb; pudb.set_trace()
 
 from context import haircontrol
 from haircontrol import discovery, data
@@ -8,36 +8,33 @@ from haircontrol import discovery, data
 class MockInspector(object):
     def __init__(self, testDataPath):
         self.testDataPath = testDataPath
+        self.e = None
 
     def connect(self, e):
-        self.name = e.name
+        self.e = e
 
     def disconnect(self):
-        self.name = None
+        self.e = None
 
     def command(self, command):
-        if not self.name:
+        if not self.e:
             return None
-        mockfile = self.testDataPath + '/' + self.name + '-' + command + '.out'
+        mockfile = self.testDataPath + '/' + self.e.name + '-' + command + '.out'
         return open(mockfile)
 
 
 class TestDiscovery(unittest.TestCase):
     
-    # http://stackoverflow.com/questions/3768895/how-to-make-a-class-json-serializable
     def setUp(self):
         self.discovery = discovery.Discovery(MockInspector('../test-data/input'))
         self.maxDiff=None
-        self.ref_equipments = [] #json.load('../test-data/ref-output/equipments.json')
 
     def test_wire_graph(self):
-        self.discovery.discover_hinting_from_lldp(data.Equipment('stg2'))
-        self.discovery.discover_from_root(data.Equipment('stg'))
-        for e in self.discovery.equipments:
-            print(e)
-            for i in e.ifaces:
-                print(i)
-        self.assertEqual(self.discovery.equipments, self.ref_equipments)
+        ref_equipments = { 'todo': data.Equipment() } #json.load('../test-data/ref-output/equipments.json')
+
+        self.discovery.discover_hinting_from_lldp(data.Equipment('stg2', '172.16.0.254'))
+        self.discovery.discover_from_root(data.Equipment('stg', '172.16.0.253'))
+        self.assertDictEqual(self.discovery.net.equipments, ref_equipments)
 
 if __name__ == '__main__':
     unittest.main()