test/lib/test_util.py

   1 #!/usr/bin/env python
   2 # vim:ts=4:sw=4:et:ai:sts=4
   3
   4 import sys
   5 import nepi.util.environ
   6
   7 # Unittest from Python 2.6 doesn't have these decorators
   8 def _bannerwrap(f, text):
   9     name = f.__name__
  10     def banner(*args, **kwargs):
  11         sys.stderr.write("*** WARNING: Skipping test %s: `%s'\n" %
  12                 (name, text))
  13         return None
  14     return banner
  15 def skip(text):
  16     return lambda f: _bannerwrap(f, text)
  17 def skipUnless(cond, text):
  18     return (lambda f: _bannerwrap(f, text)) if not cond else lambda f: f
  19 def skipIf(cond, text):
  20     return (lambda f: _bannerwrap(f, text)) if cond else lambda f: f
  21
  22 # SSH stuff
  23
  24 import os, os.path, re, signal, shutil, socket, subprocess, tempfile
  25 def gen_ssh_keypair(filename):
  26     ssh_keygen = nepi.util.environ.find_bin_or_die("ssh-keygen")
  27     args = [ssh_keygen, '-q', '-N', '', '-f', filename]
  28     assert subprocess.Popen(args).wait() == 0
  29     return filename, "%s.pub" % filename
  30
  31 def add_key_to_agent(filename):
  32     ssh_add = nepi.util.environ.find_bin_or_die("ssh-add")
  33     args = [ssh_add, filename]
  34     null = file("/dev/null", "w")
  35     assert subprocess.Popen(args, stderr = null).wait() == 0
  36     null.close()
  37
  38 def get_free_port():
  39     s = socket.socket()
  40     s.bind(("127.0.0.1", 0))
  41     port = s.getsockname()[1]
  42     return port
  43
  44 _SSH_CONF = """ListenAddress 127.0.0.1:%d
  45 Protocol 2
  46 HostKey %s
  47 UsePrivilegeSeparation no
  48 PubkeyAuthentication yes
  49 PasswordAuthentication no
  50 AuthorizedKeysFile %s
  51 UsePAM no
  52 AllowAgentForwarding yes
  53 PermitRootLogin yes
  54 StrictModes no
  55 PermitUserEnvironment yes
  56 """
  57
  58 def gen_sshd_config(filename, port, server_key, auth_keys):
  59     conf = open(filename, "w")
  60     text = _SSH_CONF % (port, server_key, auth_keys)
  61     conf.write(text)
  62     conf.close()
  63     return filename
  64
  65 def gen_auth_keys(pubkey, output, environ):
  66     #opts = ['from="127.0.0.1/32"'] # fails in stupid yans setup
  67     opts = []
  68     for k, v in environ.items():
  69         opts.append('environment="%s=%s"' % (k, v))
  70
  71     lines = file(pubkey).readlines()
  72     pubkey = lines[0].split()[0:2]
  73     out = file(output, "w")
  74     out.write("%s %s %s\n" % (",".join(opts), pubkey[0], pubkey[1]))
  75     out.close()
  76     return output
  77
  78 def start_ssh_agent():
  79     ssh_agent = nepi.util.environ.find_bin_or_die("ssh-agent")
  80     proc = subprocess.Popen([ssh_agent], stdout = subprocess.PIPE)
  81     (out, foo) = proc.communicate()
  82     assert proc.returncode == 0
  83     d = {}
  84     for l in out.split("\n"):
  85         match = re.search("^(\w+)=([^ ;]+);.*", l)
  86         if not match:
  87             continue
  88         k, v = match.groups()
  89         os.environ[k] = v
  90         d[k] = v
  91     return d
  92
  93 def stop_ssh_agent(data):
  94     # No need to gather the pid, ssh-agent knows how to kill itself; after we
  95     # had set up the environment
  96     ssh_agent = nepi.util.environ.find_bin_or_die("ssh-agent")
  97     null = file("/dev/null", "w")
  98     proc = subprocess.Popen([ssh_agent, "-k"], stdout = null)
  99     null.close()
 100     assert proc.wait() == 0
 101     for k in data:
 102         del os.environ[k]
 103
 104