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