Move stuff into a submodule. Calling it "chutney" for now. Add launcher script

3 files changed, 640 insertions, 0 deletions

diff --git a/lib/chutney/Templating.py b/lib/chutney/Templating.py
new file mode 100644
index 0000000..cc1463a
--- /dev/null
+++ b/lib/chutney/Templating.py
@@ -0,0 +1,215 @@
+#!/usr/bin/python
+#
+# Copyright 2011 Nick Mathewson, Michael Stone
+#
+#  You may do anything with this work that copyright law would normally
+#  restrict, so long as you retain the above notice(s) and this license
+#  in all redistributed copies and derived works.  There is no warranty.
+
+"""
+>>> base = Environ(foo=99, bar=600)
+>>> derived1 = Environ(parent=base, bar=700, quux=32)
+>>> base["foo"]
+99
+>>> sorted(base.keys())
+['bar', 'foo']
+>>> derived1["foo"]
+99
+>>> base["bar"]
+600
+>>> derived1["bar"]
+700
+>>> derived1["quux"]
+32
+>>> sorted(derived1.keys())
+['bar', 'foo', 'quux']
+>>> class Specialized(Environ):
+...    def __init__(self, p=None, **kw):
+...       Environ.__init__(self, p, **kw)
+...       self._n_calls = 0
+...    def _get_expensive_value(self, me):
+...       self._n_calls += 1
+...       return "Let's pretend this is hard to compute"
+...
+>>> s = Specialized(base, quux="hi")
+>>> s["quux"]
+'hi'
+>>> s['expensive_value']
+"Let's pretend this is hard to compute"
+>>> s['expensive_value']
+"Let's pretend this is hard to compute"
+>>> s._n_calls
+1
+>>> sorted(s.keys())
+['bar', 'expensive_value', 'foo', 'quux']
+
+>>> bt = _BetterTemplate("Testing ${hello}, $goodbye$$, $foo , ${a:b:c}")
+>>> bt.safe_substitute({'a:b:c': "4"}, hello=1, goodbye=2, foo=3)
+'Testing 1, 2$, 3 , 4'
+
+>>> t = Template("${include:/dev/null} $hi_there")
+>>> sorted(t.freevars())
+['hi_there']
+>>> t.format(dict(hi_there=99))
+' 99'
+>>> t2 = Template("X$${include:$fname} $bar $baz")
+>>> t2.format(dict(fname="/dev/null", bar=33, baz="$foo", foo=1337))
+'X 33 1337'
+>>> sorted(t2.freevars({'fname':"/dev/null"}))
+['bar', 'baz', 'fname']
+
+"""
+
+from __future__ import with_statement
+
+import string
+import os
+import re
+
+#class _KeyError(KeyError):
+#    pass
+
+_KeyError = KeyError
+
+class _DictWrapper:
+    def __init__(self, parent=None):
+        self._parent = parent
+
+    def __getitem__(self, key):
+        try:
+            return self._getitem(key)
+        except KeyError:
+            pass
+        if self._parent is None:
+            raise _KeyError(key)
+
+        try:
+            return self._parent[key]
+        except KeyError:
+            raise _KeyError(key)
+
+class Environ(_DictWrapper):
+    def __init__(self, parent=None, **kw):
+        _DictWrapper.__init__(self, parent)
+        self._dict = kw
+        self._cache = {}
+
+    def _getitem(self, key):
+        try:
+            return self._dict[key]
+        except KeyError:
+            pass
+        try:
+            return self._cache[key]
+        except KeyError:
+            pass
+        fn = getattr(self, "_get_%s"%key, None)
+        if fn is not None:
+            try:
+                self._cache[key] = rv = fn(self)
+                return rv
+            except _KeyError:
+                raise KeyError(key)
+        raise KeyError(key)
+
+    def __setitem__(self, key, val):
+        self._dict[key] = val
+
+    def keys(self):
+        s = set()
+        s.update(self._dict.keys())
+        s.update(self._cache.keys())
+        if self._parent is not None:
+            s.update(self._parent.keys())
+        s.update(name[5:] for name in dir(self) if name.startswith("_get_"))
+        return s
+
+class IncluderDict(_DictWrapper):
+    def __init__(self, parent, includePath=(".",)):
+        _DictWrapper.__init__(self, parent)
+        self._includePath = includePath
+        self._st_mtime = 0
+
+    def _getitem(self, key):
+        if not key.startswith("include:"):
+            raise KeyError(key)
+
+        filename = key[len("include:"):]
+        if os.path.isabs(filename):
+            with open(filename, 'r') as f:
+                stat = os.fstat(f.fileno())
+                if stat.st_mtime > self._st_mtime:
+                    self._st_mtime = stat.st_mtime
+                return f.read()
+
+        for elt in self._includePath:
+            fullname = os.path.join(elt, filename)
+            if os.path.exists(fullname):
+                with open(fullname, 'r') as f:
+                    stat = os.fstat(f.fileno())
+                    if stat.st_mtime > self._st_mtime:
+                        self._st_mtime = stat.st_mtime
+                    return f.read()
+
+        raise KeyError(key)
+
+    def getUpdateTime(self):
+        return self._st_mtime
+
+class _BetterTemplate(string.Template):
+
+    idpattern = r'[a-z0-9:_/\.\-]+'
+
+    def __init__(self, template):
+        string.Template.__init__(self, template)
+
+class _FindVarsHelper:
+    def __init__(self, dflts):
+        self._dflts = dflts
+        self._vars = set()
+    def __getitem__(self, var):
+        self._vars.add(var)
+        try:
+            return self._dflts[var]
+        except KeyError:
+            return ""
+
+class Template:
+    MAX_ITERATIONS = 32
+
+    def __init__(self, pattern, includePath=(".",)):
+        self._pat = pattern
+        self._includePath = includePath
+
+    def freevars(self, defaults=None):
+        if defaults is None:
+            defaults = {}
+        d = _FindVarsHelper(defaults)
+        self.format(d)
+        return d._vars
+
+    def format(self, values):
+        values = IncluderDict(values, self._includePath)
+        orig_val = self._pat
+        nIterations = 0
+        while True:
+            v = _BetterTemplate(orig_val).substitute(values)
+            if v == orig_val:
+                return v
+            orig_val = v
+            nIterations += 1
+            if nIterations > self.MAX_ITERATIONS:
+                raise ValueError("Too many iterations in expanding template!")
+
+if __name__ == '__main__':
+    import sys
+    if len(sys.argv) == 1:
+        import doctest
+        doctest.testmod()
+        print "done"
+    else:
+        for fn in sys.argv[1:]:
+            with open(fn, 'r') as f:
+                t = Template(f.read())
+                print fn, t.freevars()
+
diff --git a/lib/chutney/TorNet.py b/lib/chutney/TorNet.py
new file mode 100644
index 0000000..4930b17
--- /dev/null
+++ b/lib/chutney/TorNet.py
@@ -0,0 +1,423 @@
+#!/usr/bin/python
+#
+# Copyright 2011 Nick Mathewson, Michael Stone
+#
+#  You may do anything with this work that copyright law would normally
+#  restrict, so long as you retain the above notice(s) and this license
+#  in all redistributed copies and derived works.  There is no warranty.
+
+from __future__ import with_statement
+import os
+import signal
+import subprocess
+import sys
+import re
+import errno
+import time
+
+import chutney.Templating
+
+
+def mkdir_p(d):
+    try:
+        os.makedirs(d)
+    except OSError, e:
+        if e.errno == errno.EEXIST:
+            return
+        raise
+
+class Node:
+    ########
+    # Users are expected to call these:
+    def __init__(self, parent=None, **kwargs):
+        self._parent = parent
+        self._fields = self._createEnviron(parent, kwargs)
+
+    def getN(self, N):
+        return [ Node(self) for i in xrange(N) ]
+
+    def specialize(self, **kwargs):
+        return Node(parent=self, **kwargs)
+
+    #######
+    # Users are NOT expected to call these:
+
+    def _getTorrcFname(self):
+        t = chutney.Templating.Template("${torrc_fname}")
+        return t.format(self._fields)
+
+    def _createTorrcFile(self, checkOnly=False):
+        template = self._getTorrcTemplate()
+        env = self._fields
+        fn_out = self._getTorrcFname()
+        output = template.format(env)
+        if checkOnly:
+            return
+        with open(fn_out, 'w') as f:
+            f.write(output)
+
+    def _getTorrcTemplate(self):
+        env = self._fields
+        template_path = env['torrc_template_path']
+
+        t = "$${include:$torrc}"
+        return chutney.Templating.Template(t, includePath=template_path)
+
+    def _getFreeVars(self):
+        template = self._getTorrcTemplate()
+        env = self._fields
+        return template.freevars(env)
+
+    def _createEnviron(self, parent, argdict):
+        if parent:
+            parentfields = parent._fields
+        else:
+            parentfields = self._getDefaultFields()
+        return TorEnviron(parentfields, **argdict)
+
+    def _getDefaultFields(self):
+        return _BASE_FIELDS
+
+    def _checkConfig(self, net):
+        self._createTorrcFile(checkOnly=True)
+
+    def _preConfig(self, net):
+        self._makeDataDir()
+        if self._fields['authority']:
+            self._genAuthorityKey()
+        if self._fields['relay']:
+            self._genRouterKey()
+
+    def _config(self, net):
+        self._createTorrcFile()
+        #self._createScripts()
+
+    def _postConfig(self, net):
+        #self.net.addNode(self)
+        pass
+
+    def _setnodenum(self, num):
+        self._fields['nodenum'] = num
+
+    def _makeDataDir(self):
+        env = self._fields
+        datadir = env['dir']
+        mkdir_p(os.path.join(datadir, 'keys'))
+
+    def _genAuthorityKey(self):
+        env = self._fields
+        datadir = env['dir']
+        tor_gencert = env['tor_gencert']
+        lifetime = env['auth_cert_lifetime']
+        idfile   = os.path.join(datadir,'keys',"authority_identity_key")
+        skfile   = os.path.join(datadir,'keys',"authority_signing_key")
+        certfile = os.path.join(datadir,'keys',"authority_certificate")
+        addr = "%s:%s" % (env['ip'], env['dirport'])
+        passphrase = env['auth_passphrase']
+        if all(os.path.exists(f) for f in [idfile, skfile, certfile]):
+            return
+        cmdline = [
+            tor_gencert,
+            '--create-identity-key',
+            '--passphrase-fd', '0',
+            '-i', idfile,
+            '-s', skfile,
+            '-c', certfile,
+            '-m', str(lifetime),
+            '-a', addr]
+        print "Creating identity key %s for %s with %s"%(idfile,env['nick']," ".join(cmdline))
+        p = subprocess.Popen(cmdline, stdin=subprocess.PIPE)
+        p.communicate(passphrase+"\n")
+        assert p.returncode == 0 #XXXX BAD!
+
+    def _genRouterKey(self):
+        env = self._fields
+        datadir = env['dir']
+        tor = env['tor']
+        idfile = os.path.join(datadir,'keys',"identity_key")
+        cmdline = [
+            tor,
+            "--quiet",
+            "--list-fingerprint",
+            "--orport", "1",
+            "--dirserver",
+                 "xyzzy 127.0.0.1:1 ffffffffffffffffffffffffffffffffffffffff",
+            "--datadirectory", datadir ]
+        p = subprocess.Popen(cmdline, stdout=subprocess.PIPE)
+        stdout, stderr = p.communicate()
+        fingerprint = "".join(stdout.split()[1:])
+        assert re.match(r'^[A-F0-9]{40}$', fingerprint)
+        env['fingerprint'] = fingerprint
+
+    def _getDirServerLine(self):
+        env = self._fields
+        if not env['authority']:
+            return ""
+
+        datadir = env['dir']
+        certfile = os.path.join(datadir,'keys',"authority_certificate")
+        v3id = None
+        with open(certfile, 'r') as f:
+            for line in f:
+                if line.startswith("fingerprint"):
+                    v3id = line.split()[1].strip()
+                    break
+
+        assert v3id is not None
+
+        return "DirServer %s v3ident=%s orport=%s %s %s:%s %s\n" %(
+            env['nick'], v3id, env['orport'], env['dirserver_flags'],
+            env['ip'], env['dirport'], env['fingerprint'])
+
+
+    ##### Controlling a node.  This should probably get split into its
+    # own class. XXXX
+
+    def getPid(self):
+        env = self._fields
+        pidfile = os.path.join(env['dir'], 'pid')
+        if not os.path.exists(pidfile):
+            return None
+
+        with open(pidfile, 'r') as f:
+            return int(f.read())
+
+    def isRunning(self, pid=None):
+        env = self._fields
+        if pid is None:
+            pid = self.getPid()
+        if pid is None:
+            return False
+
+        try:
+            os.kill(pid, 0) # "kill 0" == "are you there?"
+        except OSError, e:
+            if e.errno == errno.ESRCH:
+                return False
+            raise
+
+	# okay, so the process exists.  Say "True" for now.
+        # XXXX check if this is really tor!
+        return True
+
+    def check(self, listRunning=True, listNonRunning=False):
+        env = self._fields
+        pid = self.getPid()
+        running = self.isRunning(pid)
+        name = env['nick']
+        dir = env['dir']
+        if running:
+            if listRunning:
+                print "%s is running with PID %s"%(name,pid)
+            return True
+        elif os.path.exists(os.path.join(dir, "core.%s"%pid)):
+            if listNonRunning:
+                print "%s seems to have crashed, and left core file core.%s"%(
+                   nick,pid)
+            return False
+        else:
+            if listNonRunning:
+                print "%s is stopped"%nick
+            return False
+
+    def hup(self):
+        pid = self.getPid()
+        running = self.isRunning()
+        nick = self._fields['nick']
+        if self.isRunning():
+            print "Sending sighup to %s"%nick
+            os.kill(pid, signal.SIGHUP)
+            return True
+        else:
+            print "%s is not running"%nick
+            return False
+
+    def start(self):
+        if self.isRunning():
+            print "%s is already running"
+            return
+        torrc = self._getTorrcFname()
+        cmdline = [
+            self._fields['tor'],
+            "--quiet",
+            "-f", torrc,
+            ]
+        p = subprocess.Popen(cmdline)
+        # XXXX this requires that RunAsDaemon is set.
+        p.wait()
+        if p.returncode != 0:
+            print "Couldn't launch %s (%s): %s"%(self._fields['nick'],
+                                                 " ".join(cmdline),
+                                                 p.returncode)
+            return False
+        return True
+
+    def stop(self, sig=signal.SIGINT):
+        env = self._fields
+        pid = self.getPid()
+        if not self.isRunning(pid):
+            print "%s is not running"%env['nick']
+            return
+        os.kill(pid, sig)
+
+
+DEFAULTS = {
+    'authority' : False,
+    'relay' : False,
+    'connlimit' : 60,
+    'net_base_dir' : 'net',
+    'tor' : 'tor',
+    'auth_cert_lifetime' : 12,
+    'ip' : '127.0.0.1',
+    'dirserver_flags' : 'no-v2',
+    'privnet_dir' : '.',
+    'torrc_fname' : '${dir}/torrc',
+    'orport_base' : 6000,
+    'dirport_base' : 7000,
+    'controlport_base' : 8000,
+    'socksport_base' : 9000,
+    'dirservers' : "Dirserver bleargh bad torrc file!",
+    'core' : True,
+}
+
+class TorEnviron(chutney.Templating.Environ):
+    def __init__(self,parent=None,**kwargs):
+        chutney.Templating.Environ.__init__(self, parent=parent, **kwargs)
+
+    def _get_orport(self, me):
+        return me['orport_base']+me['nodenum']
+
+    def _get_controlport(self, me):
+        return me['controlport_base']+me['nodenum']
+
+    def _get_socksport(self, me):
+        return me['socksport_base']+me['nodenum']
+
+    def _get_dirport(self, me):
+        return me['dirport_base']+me['nodenum']
+
+    def _get_dir(self, me):
+        return os.path.abspath(os.path.join(me['net_base_dir'],
+                                            "nodes",
+                                         "%03d%s"%(me['nodenum'], me['tag'])))
+
+    def _get_nick(self, me):
+        return "test%03d%s"%(me['nodenum'], me['tag'])
+
+    def _get_tor_gencert(self, me):
+        return me['tor']+"-gencert"
+
+    def _get_auth_passphrase(self, me):
+        return self['nick'] # OMG TEH SECURE!
+
+    def _get_torrc_template_path(self, me):
+        return [ os.path.join(me['privnet_dir'], 'torrc_templates') ]
+
+
+class Network:
+    def __init__(self,defaultEnviron):
+        self._nodes = []
+        self._dfltEnv = defaultEnviron
+        self._nextnodenum = 0
+
+    def _addNode(self, n):
+        n._setnodenum(self._nextnodenum)
+        self._nextnodenum += 1
+        self._nodes.append(n)
+
+    def _checkConfig(self):
+        for n in self._nodes:
+            n._checkConfig(self)
+
+    def configure(self):
+        network = self
+        dirserverlines = []
+
+        self._checkConfig()
+
+        # XXX don't change node names or types or count if anything is
+        # XXX running!
+
+        for n in self._nodes:
+            n._preConfig(network)
+            dirserverlines.append(n._getDirServerLine())
+
+        self._dfltEnv['dirservers'] = "".join(dirserverlines)
+
+        for n in self._nodes:
+            n._config(network)
+
+        for n in self._nodes:
+            n._postConfig(network)
+
+    def status(self):
+        statuses = [n.check() for n in self._nodes]
+        n_ok = len([x for x in statuses if x])
+        print "%d/%d nodes are running"%(n_ok,len(self._nodes))
+
+    def restart(self):
+        self.stop()
+        self.start()
+
+    def start(self):
+        print "Starting nodes"
+        return all([n.start() for n in self._nodes])
+
+    def hup(self):
+        print "Sending SIGHUP to nodes"
+        return all([n.hup() for n in self._nodes])
+
+    def stop(self):
+        for sig, desc in [(signal.SIGINT, "SIGINT"),
+                          (signal.SIGINT, "another SIGINT"),
+                          (signal.SIGKILL, "SIGKILL")]:
+            print "Sending %s to nodes"%desc
+            for n in self._nodes:
+                if n.isRunning():
+                    n.stop(sig=sig)
+            print "Waiting for nodes to finish."
+            for n in xrange(15):
+                time.sleep(1)
+                if all(not n.isRunning() for n in self._nodes):
+                    return
+                sys.stdout.write(".")
+                sys.stdout.flush()
+            for n in self._nodes:
+                n.check(listNonRunning=False)
+
+def ConfigureNodes(nodelist):
+    network = _THE_NETWORK
+
+    for n in nodelist:
+        network._addNode(n)
+
+def runConfigFile(verb, f):
+    global _BASE_FIELDS
+    global _THE_NETWORK
+    _BASE_FIELDS = TorEnviron(chutney.Templating.Environ(**DEFAULTS))
+    _THE_NETWORK = Network(_BASE_FIELDS)
+
+    _GLOBALS = dict(_BASE_FIELDS= _BASE_FIELDS,
+                    Node=Node,
+                    ConfigureNodes=ConfigureNodes,
+                    _THE_NETWORK=_THE_NETWORK)
+
+    exec f in _GLOBALS
+    network = _GLOBALS['_THE_NETWORK']
+
+    if not hasattr(network, verb):
+        print "I don't know how to %s.  Known commands are: %s" % (
+            verb, " ".join(x for x in dir(network) if not x.startswith("_")))
+        return
+
+    getattr(network,verb)()
+
+if __name__ == '__main__':
+    if len(sys.argv) < 3:
+        print "Syntax: chutney {command} {networkfile}"
+        sys.exit(1)
+
+    f = open(sys.argv[2])
+    runConfigFile(sys.argv[1], f)
+
+
diff --git a/lib/chutney/__init__.py b/lib/chutney/__init__.py
new file mode 100644
index 0000000..139597f
--- /dev/null
+++ b/lib/chutney/__init__.py
@@ -0,0 +1,2 @@
+
+