Refactoring, dicover_from_root -> _from_gateway. Also calls compute_neighbourhood.HEADmaster

3 files changed, 47 insertions, 25 deletions

diff --git a/haircontrol/data.py b/haircontrol/data.py
index 1eb4d3b..da85100 100644
--- a/haircontrol/data.py
+++ b/haircontrol/data.py
@@ -3,10 +3,26 @@ class EtherDomain:
         self.equipments = {}
         self.ip2mac = {}
         self.mac2ip = {}
+        self.gw = None
+        self.root_iface = None
 
     def __repr__(self):
         return '([\n%s\n],\n%s\n)'%(',\n'.join('  %s'%repr(e) for e in self.get_equipment_list_sorted()), repr(self.ip2mac))
 
+    def print_tree(self):
+        self._print_tree(self.gw, self.root_iface, 0)
+
+    def _print_tree(self,eq_root, iface, indent):
+        print(" "*indent + eq_root.name + '(' + iface.name + ')')
+        ei_pairs = []
+        for iface in eq_root.ifaces.values():
+            for ip in iface.direct_neighbours:
+                ei_pairs.append( (self.equipments[ip], iface) )
+        ei_pairs.sort(key=lambda x: x[0].name)
+
+        for (e,i) in ei_pairs:
+            self._print_tree(e, i, indent+1)
+
     def get_equipment_list_sorted(self):
         e_list = list(self.equipments.values())
         e_list.sort(key=lambda x: x.name)
@@ -29,11 +45,13 @@ class EtherDomain:
         self.mac2ip[mac] = ip
 
 class Equipment:
-    def __init__(self, name=None, mgmtip=None):
+    def __init__(self, name=None, mgmtip=None, first_ifname=None):
         self.name = name
         self.mgmtip = mgmtip
         self.ifaces = {}
         self.inspected = False
+        if first_ifname:
+            self.ifaces[first_ifname] = Interface(first_ifname, None)
 
     def __repr__(self):
         return repr( (self.name, self.mgmtip, list(self.ifaces.values())) )
diff --git a/haircontrol/discovery.py b/haircontrol/discovery.py
index 9fd29da..c072615 100644
--- a/haircontrol/discovery.py
+++ b/haircontrol/discovery.py
@@ -6,6 +6,8 @@ class Discovery:
 
     def __init__(self): #, inspector):
         self.net = EtherDomain()
+        # XXX make that configurable
+        self.ignore_ip = ['172.16.0.253']
         self.dummy_inspector = DummyInspector()
         self.linux_inspector = LinuxInspector()
         self.ubnt_inspector = UbntInspector()
@@ -99,21 +101,29 @@ class Discovery:
         self.linux_inspector.disconnect()
 
 
-    def discover_from_root(self, e_root):
-        self.net.add_equipment(e_root)
-        self.linux_inspector.connect(e_root)
+    def discover_from_gateway(self, e_gw):
+        self.net.gw = e_gw
+        if e_gw.ifaces:
+            self.net.root_iface = list(e_gw.ifaces.values())[0]
+        self.net.add_equipment(e_gw)
+        self._learn_from_gw_neigh()
+        self._inspect_all()
+        self._compute_neighbourhood()
 
-        # Learn root neighbours (directly or indirectly connected via trasparent bridges)
+    # Learn gateway neighbours (directly or indirectly connected via trasparent bridges)
+    def _learn_from_gw_neigh(self):
+        e_gw = self.net.gw
+        self.linux_inspector.connect(e_gw)
         result = self.linux_inspector.command('ip-neigh')
         for (ip, ifname, mac) in result:
             self.net.index_mac_ip(mac, ip)
-            e_root.add_seen_mac(ifname, mac)
+            e_gw.add_seen_mac(ifname, mac)
 
         self.linux_inspector.disconnect()
 
         # Create/Update Equipment object for all neighbours
         # (could be already created by hinting)
-        for iface in e_root.ifaces.values():
+        for iface in e_gw.ifaces.values():
             local_ifname = iface.name
             local_mac = iface.mac
             for remote_mac in iface.mac_seen:
@@ -128,6 +138,7 @@ class Discovery:
                         self.net.add_equipment(e)
 
         
+    def _inspect_all(self):
         # Inspect all non-already inspected equipement
         done = False
         while not done:
@@ -191,18 +202,9 @@ class Discovery:
                         print("Notice: Nothing inspected on %s"%e)
                     i.disconnect()
 
-    def compute_neighbourhood(self):
+    def _compute_neighbourhood(self):
 
-        ### Configuration
-
-        # GATEWAY_IP and GATEWAY_IFACE_NAME define information about
-        # the gateway, i.e., the center of the network.
-        #
-        # IGNORE_IP lists all IP that should not appear in the network
-        # representation.
-        gateway_ip = '172.16.0.254'
-        gateway_iface_name = 'eth1'
-        ignore_ip = ['172.16.0.253']
+        gateway_ip = self.net.gw.mgmtip
 
         ### State variables
 
@@ -214,7 +216,7 @@ class Discovery:
             for interface in equipment.ifaces.values():
                 mac_to_ip[interface.mac] = ip
         equipments = { ip: eq for ip, eq in self.net.equipments.items()
-                       if (ip not in ignore_ip) and (ip != gateway_ip) }
+                       if (ip not in self.ignore_ip) and (ip != gateway_ip) }
         outer_net = []          # IP belonging to external levels.
         with_uplink = []        # IP already attached to an interface.
 
@@ -287,6 +289,8 @@ class Discovery:
                     link_to_outer_net(equipment)
                     current_level.append(ip)
             outer_net.extend(current_level)
-            for ip in current_level: del equipments[ip]
-        exit_iface = self.net.equipments[gateway_ip].ifaces[gateway_iface_name]
-        for ip in outer_net: set_uplink(ip, exit_iface)
+            for ip in current_level:
+                del equipments[ip]
+        for ip in outer_net:
+            set_uplink(ip, self.net.root_iface)
+
diff --git a/tests/test_discovery.py b/tests/test_discovery.py
index d89d6f6..35600ae 100755
--- a/tests/test_discovery.py
+++ b/tests/test_discovery.py
@@ -19,10 +19,10 @@ class TestDiscovery(unittest.TestCase):
             ('SW_PI_EGL', '172.16.30.27'),
             ('SW_Eglise_ESTANCARBON', '172.16.30.38'),
         ])
-        self.discovery.discover_from_root(data.Equipment('stg', '172.16.0.254'))
-        self.discovery.compute_neighbourhood()
+        self.discovery.discover_from_gateway(data.Equipment('stg', '172.16.0.254', 'eth1'))
 #        self.assertEqual(self.ref_net.get_equipment_list_sorted(), list(self.discovery.net.get_equipment_list_sorted()))
-        print(self.discovery.net) # <----- ya un gros print ici qui affiche... tout.
+#        print(self.discovery.net) # <----- ya un gros print ici qui affiche... tout.
+        self.discovery.net.print_tree()
 
 if __name__ == '__main__':
     unittest.main()