demo: some clean-up in hmac test

README.rst

@@ -58,7 +58,7 @@ Requirements
 - Babel_ (with Nexedi patches)
 - geoip2: `python library`_ and `country lite database`_ (optional)
 - python-miniupnpc for UPnP support (optional)
-- for the demo: miniupnpd_, Graphviz, Screen_, Nemu_, MultiPing_
+- for the demo: miniupnpd_, Graphviz, Screen_, Nemu_, MultiPing_, psutil_
 
 See also `setup.py` for Python dependencies.
 
@@ -66,6 +66,7 @@ See also `setup.py` for Python dependencies.
 .. _Nemu: https://github.com/thetincho/nemu
 .. _miniupnpd: http://miniupnp.free.fr/
 .. _MultiPing: https://github.com/romana/multi-ping
+.. _psutil: https://pypi.org/project/psutil/
 .. _Screen: http://savannah.gnu.org/projects/screen
 .. _python library: https://pypi.org/project/geoip2/
 .. _country lite database: https://dev.maxmind.com/geoip/geoip2/geolite2/

demo/demo

@@ -246,8 +246,7 @@ def new_network(registry, reg_addr, serial, ca):
         """)).wait()
     db = sqlite3.connect(db_path, isolation_level=None)
     def new_node(node, folder, args='', prefix_len=None, registry=registry_url):
-        if node not in nodes:
-            nodes.append(node)
+        nodes.append(node)
         if not os.path.exists(folder + '/cert.crt'):
             dh_path = folder + '/dh2048.pem'
             if not os.path.exists(dh_path):
@@ -267,9 +266,11 @@ def new_network(registry, reg_addr, serial, ca):
             p.communicate(str(token[0]))
             os.remove(dh_path)
             os.remove(folder + '/ca.crt')
-        node.screen('./py re6stnet @%s/re6stnet.conf -v%u --registry %s'
-                    ' --console %s/run/console.sock %s' % (
-                      folder, VERBOSE, registry, folder, args))
+        node.re6st_cmdline = (
+            './py re6stnet @%s/re6stnet.conf -v%u --registry %s'
+            ' --console %s/run/console.sock %s'
+            ) % (folder, VERBOSE, registry, folder, args)
+        node.screen(node.re6st_cmdline)
     new_node(registry, registry.name, '--ip ' + reg_addr, registry='http://localhost/')
     yield new_node
     db.close()
@@ -304,42 +305,49 @@ if args.ping:
 class testHMAC(Thread):
 
     def run(self):
+        updateHMAC = ('python', '-c', "import urllib, sys; sys.exit("
+            "204 != urllib.urlopen('http://127.0.0.1/updateHMAC').code)")
         reg1_db = sqlite3.connect('registry/registry.db', isolation_level=None,
                                                  check_same_thread=False)
         reg2_db = sqlite3.connect('registry2/registry.db', isolation_level=None,
                                                  check_same_thread=False)
         reg1_db.text_factory = reg2_db.text_factory = str
-        m_net1 = ['registry', 'm1', 'm2', 'm3', 'm4', 'm5', 'm6', 'm7', 'm8']
-        m_net2 = ['registry2', 'm10']
+        m_net1 = 'registry', 'm1', 'm2', 'm3', 'm4', 'm5', 'm6', 'm7', 'm8'
+        m_net2 = 'registry2', 'm10'
         print 'Testing HMAC, letting the time to machines to create tunnels...'
         time.sleep(45)
         print 'Check that the initial HMAC config is deployed on network 1'
         test_hmac.checkHMAC(reg1_db, m_net1)
         print 'Test that a HMAC update works with nodes that are up'
-        registry.screen('wget http://10.0.0.2/updateHMAC')
+        registry.backticks_raise(updateHMAC)
         print 'Updated HMAC (config = hmac0 & hmac1), waiting...'
         time.sleep(60)
         print 'Checking HMAC on machines connected to registry 1...'
         test_hmac.checkHMAC(reg1_db, m_net1)
         print ('Test that machines can update upon reboot ' +
                'when they were off during a HMAC update.')
-        test_hmac.killRe6st('m1')
+        test_hmac.killRe6st(machine1)
         print 'Re6st on machine 1 is stopped'
         time.sleep(5)
-        registry.screen('wget http://10.0.0.2/updateHMAC')
+        registry.backticks_raise(updateHMAC)
         print 'Updated HMAC on registry (config = hmac1 & hmac2), waiting...'
         time.sleep(60)
-        new_node(machine1, 'm1', '-I%s' % m1_if_0.name,
-                 None, 'http://%s/' % REGISTRY)
+        machine1.screen(machine1.re6st_cmdline)
         print 'Started re6st on machine 1, waiting for it to get new conf'
         time.sleep(60)
         print 'Checking HMAC on machines connected to registry 1...'
         test_hmac.checkHMAC(reg1_db, m_net1)
         print 'Testing of HMAC done!'
+        # TODO: missing last step
+        reg1_db.close()
+        reg2_db.close()
 
 if args.hmac:
     import test_hmac
-    testHMAC().start()
+    t = testHMAC()
+    t.deamon = 1
+    t.start()
+    del t
 
 _ll = {}
 def node_by_ll(addr):

demo/test_hmac.py

-import sqlite3, subprocess
+from binascii import b2a_hex
+import psutil
+
+BABEL_HMAC = 'babel_hmac0', 'babel_hmac1', 'babel_hmac2'
 
 def getConfig(db, name):
-    r, = next(db.execute(
-              "SELECT value FROM config WHERE name=?", (name,)), (None,))
-    if r is not None:
-        r = str(r).encode('hex')
-    return r
+    r = db.execute("SELECT value FROM config WHERE name=?", (name,)).fetchone()
+    if r:
+        return b2a_hex(*r)
 
-def killRe6st(machine):
-    p = subprocess.Popen(['pgrep', '-f', 'set ./py re6stnet @%s' %machine],
-                         stdout=subprocess.PIPE)
-    ps_id = p.communicate()[0].split('\n', 1)[0]
-    if ps_id:
-        subprocess.Popen(['kill', ps_id])
+def killRe6st(node):
+    for p in psutil.Process(node._screen.pid).children():
+        if p.cmdline()[-1].startswith('set ./py re6stnet '):
+            p.kill()
+            break
 
 def checkHMAC(db, machines):
-    hmac = dict([(k, getConfig(db, k))
-                for k in 'babel_hmac0', 'babel_hmac1', 'babel_hmac2'])
+    hmac = [getConfig(db, k) for k in BABEL_HMAC]
     rc = True
-    ps = subprocess.Popen(['pgrep', '-a', 'babel'], stdout=subprocess.PIPE)
-    for p in (p for p in ps.communicate()[0].split('\n') if p):
-        if p.split('/',1)[0].split()[-1] in machines:
-            if hmac['babel_hmac0'] and not hmac['babel_hmac1']: # state = hmac0
-                if ('sign' not in p or
-                    'accept' in p or
-                    p.split('sign value ',1)[1].split()[0]\
-                      != hmac['babel_hmac0']):
-                    rc = False
-                    print 'HMAC config wrong for in %s' % p
+    for x in psutil.Process().children(True):
+        if x.name() == 'babeld':
+            sign = accept = None
+            args = x.cmdline()
+            for x in args:
+                if x.endswith('/babeld.log'):
+                    if x[:-11] not in machines:
+                        break
+                elif x.startswith('key '):
+                    x = x.split()
+                    if 'sign' in x:
+                        sign = x[-1]
+                    elif 'accept' in x:
+                        accept = x[-1]
             else:
-                if hmac['babel_hmac0']: # state = hmac0 and hmac1
-                    sign = 'babel_hmac0'
-                    accept = 'babel_hmac1'
-                else: # state = hmac1 and hmac2
-                    sign = 'babel_hmac1'
-                    accept = 'babel_hmac2'
-                if ('accept' not in p or
-                    'sign' not in p or
-                    p.split('sign value ',1)[1].split()[0] != hmac[sign] or
-                    p.split('accept value ',1)[1].split()[0] != hmac[accept]):
+                i = 0 if hmac[0] else 1
+                if hmac[i] != sign or hmac[i+1] != accept:
+                    print 'HMAC config wrong for in %s' % args
                     rc = False
-                    print 'HMAC config wrong in %s' % p
     if rc:
         print('All nodes use Babel with the correct HMAC configuration')
     else:
-        print('Correct config: %s' % hmac)
+        print('Expected config: %s' % dict(zip(BABEL_HMAC, hmac)))
     return rc