src/nepi/util/environ.py

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