src/nepi/util/environ.py

   1 import ctypes
   2 import imp
   3 import sys
   4
   5 import os, os.path, re, signal, shutil, socket, subprocess, tempfile
   6
   7 __all__ =  ["python", "ssh_path"]
   8 __all__ += ["rsh", "tcpdump_path", "sshd_path"]
   9 __all__ += ["execute", "backticks"]
  10
  11
  12 # Unittest from Python 2.6 doesn't have these decorators
  13 def _bannerwrap(f, text):
  14     name = f.__name__
  15     def banner(*args, **kwargs):
  16         sys.stderr.write("*** WARNING: Skipping test %s: `%s'\n" %
  17                 (name, text))
  18         return None
  19     return banner
  20
  21 def skip(text):
  22     return lambda f: _bannerwrap(f, text)
  23
  24 def skipUnless(cond, text):
  25     return (lambda f: _bannerwrap(f, text)) if not cond else lambda f: f
  26
  27 def skipIf(cond, text):
  28     return (lambda f: _bannerwrap(f, text)) if cond else lambda f: f
  29
  30 def find_bin(name, extra_path = None):
  31     search = []
  32     if "PATH" in os.environ:
  33         search += os.environ["PATH"].split(":")
  34     for pref in ("/", "/usr/", "/usr/local/"):
  35         for d in ("bin", "sbin"):
  36             search.append(pref + d)
  37     if extra_path:
  38         search += extra_path
  39
  40     for d in search:
  41             try:
  42                 os.stat(d + "/" + name)
  43                 return d + "/" + name
  44             except OSError, e:
  45                 if e.errno != os.errno.ENOENT:
  46                     raise
  47     return None
  48
  49 def find_bin_or_die(name, extra_path = None):
  50     r = find_bin(name)
  51     if not r:
  52         raise RuntimeError(("Cannot find `%s' command, impossible to " +
  53                 "continue.") % name)
  54     return r
  55
  56 def find_bin(name, extra_path = None):
  57     search = []
  58     if "PATH" in os.environ:
  59         search += os.environ["PATH"].split(":")
  60     for pref in ("/", "/usr/", "/usr/local/"):
  61         for d in ("bin", "sbin"):
  62             search.append(pref + d)
  63     if extra_path:
  64         search += extra_path
  65
  66     for d in search:
  67             try:
  68                 os.stat(d + "/" + name)
  69                 return d + "/" + name
  70             except OSError, e:
  71                 if e.errno != os.errno.ENOENT:
  72                     raise
  73     return None
  74
  75 ssh_path = find_bin_or_die("ssh")
  76 python_path = find_bin_or_die("python")
  77
  78 # Optional tools
  79 rsh_path = find_bin("rsh")
  80 tcpdump_path = find_bin("tcpdump")
  81 sshd_path = find_bin("sshd")
  82
  83 def execute(cmd):
  84     # FIXME: create a global debug variable
  85     #print "[pid %d]" % os.getpid(), " ".join(cmd)
  86     null = open("/dev/null", "r+")
  87     p = subprocess.Popen(cmd, stdout = null, stderr = subprocess.PIPE)
  88     out, err = p.communicate()
  89     if p.returncode != 0:
  90         raise RuntimeError("Error executing `%s': %s" % (" ".join(cmd), err))
  91
  92 def backticks(cmd):
  93     p = subprocess.Popen(cmd, stdout = subprocess.PIPE,
  94             stderr = subprocess.PIPE)
  95     out, err = p.communicate()
  96     if p.returncode != 0:
  97         raise RuntimeError("Error executing `%s': %s" % (" ".join(cmd), err))
  98     return out
  99
 100
 101 # SSH stuff
 102
 103 def gen_ssh_keypair(filename):
 104     ssh_keygen = nepi.util.environ.find_bin_or_die("ssh-keygen")
 105     args = [ssh_keygen, '-q', '-N', '', '-f', filename]
 106     assert subprocess.Popen(args).wait() == 0
 107     return filename, "%s.pub" % filename
 108
 109 def add_key_to_agent(filename):
 110     ssh_add = nepi.util.environ.find_bin_or_die("ssh-add")
 111     args = [ssh_add, filename]
 112     null = file("/dev/null", "w")
 113     assert subprocess.Popen(args, stderr = null).wait() == 0
 114     null.close()
 115
 116 def get_free_port():
 117     s = socket.socket()
 118     s.bind(("127.0.0.1", 0))
 119     port = s.getsockname()[1]
 120     return port
 121
 122 _SSH_CONF = """ListenAddress 127.0.0.1:%d
 123 Protocol 2
 124 HostKey %s
 125 UsePrivilegeSeparation no
 126 PubkeyAuthentication yes
 127 PasswordAuthentication no
 128 AuthorizedKeysFile %s
 129 UsePAM no
 130 AllowAgentForwarding yes
 131 PermitRootLogin yes
 132 StrictModes no
 133 PermitUserEnvironment yes
 134 """
 135
 136 def gen_sshd_config(filename, port, server_key, auth_keys):
 137     conf = open(filename, "w")
 138     text = _SSH_CONF % (port, server_key, auth_keys)
 139     conf.write(text)
 140     conf.close()
 141     return filename
 142
 143 def gen_auth_keys(pubkey, output, environ):
 144     #opts = ['from="127.0.0.1/32"'] # fails in stupid yans setup
 145     opts = []
 146     for k, v in environ.items():
 147         opts.append('environment="%s=%s"' % (k, v))
 148
 149     lines = file(pubkey).readlines()
 150     pubkey = lines[0].split()[0:2]
 151     out = file(output, "w")
 152     out.write("%s %s %s\n" % (",".join(opts), pubkey[0], pubkey[1]))
 153     out.close()
 154     return output
 155
 156 def start_ssh_agent():
 157     ssh_agent = nepi.util.environ.find_bin_or_die("ssh-agent")
 158     proc = subprocess.Popen([ssh_agent], stdout = subprocess.PIPE)
 159     (out, foo) = proc.communicate()
 160     assert proc.returncode == 0
 161     d = {}
 162     for l in out.split("\n"):
 163         match = re.search("^(\w+)=([^ ;]+);.*", l)
 164         if not match:
 165             continue
 166         k, v = match.groups()
 167         os.environ[k] = v
 168         d[k] = v
 169     return d
 170
 171 def stop_ssh_agent(data):
 172     # No need to gather the pid, ssh-agent knows how to kill itself; after we
 173     # had set up the environment
 174     ssh_agent = nepi.util.environ.find_bin_or_die("ssh-agent")
 175     null = file("/dev/null", "w")
 176     proc = subprocess.Popen([ssh_agent, "-k"], stdout = null)
 177     null.close()
 178     assert proc.wait() == 0
 179     for k in data:
 180         del os.environ[k]
 181