TunnelManager is now a class

TODO

@@ -7,8 +7,6 @@ To be done :
 To be discuss:
     Remove the --no-boot option since we know when no node is avalaible
 
-    Find a better solution for config than utils.config = config, openv.config = config, ...
-    
     When I created PeersDB, I thought only be used to access the DB and not do some logic. 
         We should decide what it is suppose to do :
             Just access the DB

tunnelmanager.py

-import os, random
+import os, random, traceback
 import openvpn
 import utils
 import db
 
-connection_dict = {} # to remember current connections we made
 free_interface_set = set(('client1', 'client2', 'client3', 'client4', 'client5',
                           'client6', 'client7', 'client8', 'client9', 'client10'))
 
-def startNewConnections(n, write_pipe):
-    try:
-        for peer_id, ip, port, proto in peers_db.getUnusedPeers(n):
-            utils.log('Establishing a connection with id %s (%s:%s)' % (peer_id, ip, port), 2)
-            iface = free_interface_set.pop()
-            connection_dict[peer_id] = ( openvpn.client( ip, write_pipe, '--dev', iface, '--proto', proto, '--rport', str(port),
-                stdout=os.open(os.path.join(utils.config.log, 'vifibnet.client.%s.log' % (peer_id,)), 
-                               os.O_WRONLY|os.O_CREAT|os.O_TRUNC) ),
-                iface)
-            peers_db.usePeer(peer_id)
-    except KeyError:
-        utils.log("Can't establish connection with %s : no available interface" % ip, 2)
-    except Exception:
-        traceback.print_exc()
-
-def killConnection(peer_id):
-    try:
+class TunnelManager:
+
+    def __init__(self, write_pipe, peers_db):
+        self.write_pipe = write_pipe
+        self.peers_db = peers_db
+        self.connection_dict = {}
+
+    def refresh(self):
+        self.cleanDeads()
+        self.removeSomeTunnels()
+        self.makeNewTunnels()
+
+    def cleanDeads(self):
+        for id in self.connection_dict.keys():
+            p, iface = self.connection_dict[id]
+            if p.poll() != None:
+                utils.log('Connection with %s has failed with return code %s' % (id, p.returncode), 3)
+                free_interface_set.add(iface)
+                self.peers_db.unusePeer(id)
+                del self.connection_dict[id]
+
+    def removeSomeTunnels(self):
+        for i in range(0, max(0, len(self.connection_dict) - utils.config.client_count + utils.config.refresh_count)):
+            peer_id = random.choice(self.connection_dict.keys())
+            kill(peer_id)
+
+    def kill(self, peer_id):
         utils.log('Killing the connection with id ' + str(peer_id), 2)
-        p, iface = connection_dict.pop(peer_id)
+        p, iface = self.connection_dict.pop(peer_id)
         p.kill()
         free_interface_set.add(iface)
-        peers_db.unusePeer(peer_id)
-    except KeyError:
-        utils.log("Can't kill connection to " + peer_id + ": no existing connection", 1)
-        pass
-    except Exception:
-        utils.log("Can't kill connection to " + peer_id + ": uncaught error", 1)
-        pass
-
-def checkConnections():
-    for id in connection_dict.keys():
-        p, iface = connection_dict[id]
-        if p.poll() != None:
-            utils.log('Connection with %s has failed with return code %s' % (id, p.returncode), 3)
-            free_interface_set.add(iface)
-            peers_db.unusePeer(id)
-            del connection_dict[id]
-
-def refreshConnections(write_pipe):
-    checkConnections()
-    # Kill some random connections
-    try:
-        for i in range(0, max(0, len(connection_dict) - utils.config.client_count + utils.config.refresh_count)):
-            peer_id = random.choice(connection_dict.keys())
-            killConnection(peer_id)
-    except Exception:
-        pass
-    # Establish new connections
-    startNewConnections(utils.config.client_count - len(connection_dict), write_pipe)
+        self.peers_db.unusePeer(peer_id)
 
+    def makeNewTunnels(self):
+        try:
+            for peer_id, ip, port, proto in self.peers_db.getUnusedPeers(utils.config.client_count - len(self.connection_dict), self.write_pipe):
+                utils.log('Establishing a connection with id %s (%s:%s)' % (peer_id, ip, port), 2)
+                iface = free_interface_set.pop()
+                self.connection_dict[peer_id] = ( openvpn.client( ip, write_pipe, '--dev', iface, '--proto', proto, '--rport', str(port),
+                    stdout=os.open(os.path.join(utils.config.log, 'vifibnet.client.%s.log' % (peer_id,)), 
+                                   os.O_WRONLY|os.O_CREAT|os.O_TRUNC) ),
+                                   iface)
+                self.peers_db.usePeer(peer_id)
+        except KeyError:
+            utils.log("Can't establish connection with %s : no available interface" % ip, 2)
+        except Exception:
+            traceback.print_exc()

vifibnet.py

@@ -47,9 +47,6 @@ def main():
     # Get arguments
     utils.getConfig()
     
-    # Setup database
-    tunnelmanager.peers_db = db.PeersDB(utils.config.db)
-
     # Launch babel on all interfaces. WARNING : you have to be root to start babeld
     utils.log('Starting babel', 3)
     babel = startBabel(stdout=os.open(os.path.join(utils.config.log, 'vifibnet.babeld.log'), 
@@ -60,11 +57,15 @@ def main():
     r_pipe, write_pipe = os.pipe()
     read_pipe = os.fdopen(r_pipe)
 
-    # Establish connections
+    # setup the tunnel manager 
+    peers_db = db.PeersDB(utils.config.db)
+    tunnelManager = tunnelmanager.TunnelManager(write_pipe, peers_db)
+
+   # Establish connections
     utils.log('Starting openvpn server', 3)
     serverProcess = openvpn.server(utils.config.internal_ip, write_pipe, '--dev', 'vifibnet',
             stdout=os.open(os.path.join(utils.config.log, 'vifibnet.server.log'), os.O_WRONLY | os.O_CREAT | os.O_TRUNC))
-    tunnelmanager.startNewConnections(utils.config.client_count, write_pipe)
+    tunnelManager.refresh()
 
     # Timed refresh initializing
     next_refresh = time.time() + utils.config.refresh_time
@@ -77,8 +78,8 @@ def main():
             if ready:
                 handle_message(read_pipe.readline())
             if time.time() >= next_refresh:
-                tunnelmanager.peers_db.populate(10)
-                tunnelmanager.refreshConnections(write_pipe)
+                peers_db.populate(10)
+                tunnelManager.refresh()
                 next_refresh = time.time() + utils.config.refresh_time
     except KeyboardInterrupt:
         return 0