cmd_env.py

   1 # Copyright (C) Cmed Ltd, 2008, 2009
   2 #
   3 # This program is free software; you can redistribute it and/or modify
   4 # it under the terms of the GNU Lesser General Public License as
   5 # published by the Free Software Foundation; either version 2.1 of the
   6 # License, or (at your option) any later version.
   7 #
   8 # This program is distributed in the hope that it will be useful, but
   9 # WITHOUT ANY WARRANTY; without even the implied warranty of
  10 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  11 # Lesser General Public License for more details.
  12 #
  13 # You should have received a copy of the GNU Lesser General Public
  14 # License along with this program; if not, write to the Free Software
  15 # Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA
  16 # 02110-1301, USA.
  17
  18 import errno
  19 import os
  20 import pprint
  21 import signal
  22 import subprocess
  23 import time
  24
  25
  26 def read_file(filename):
  27     fh = open(filename, "r")
  28     try:
  29         return fh.read()
  30     finally:
  31         fh.close()
  32
  33
  34 def write_file(filename, data):
  35     fh = open(filename, "w")
  36     try:
  37         fh.write(data)
  38     finally:
  39         fh.close()
  40
  41
  42 class CommandFailedError(Exception):
  43
  44     def __init__(self, message, rc):
  45         Exception.__init__(self, message)
  46         self.rc = rc
  47
  48
  49 class BasicEnv(object):
  50
  51     def cmd(self, args, do_wait=True, fork=True, **kwargs):
  52         if fork:
  53             process = subprocess.Popen(args, **kwargs)
  54             if do_wait:
  55                 rc = process.wait()
  56                 if rc != 0:
  57                     raise CommandFailedError(
  58                         "Command failed with return code %i: %s" % (rc, args),
  59                         rc)
  60             return process
  61         else:
  62             # TODO: support other keyword args
  63             assert len(kwargs) == 0, kwargs
  64             os.execvp(args[0], args)
  65
  66
  67 def call(args, **kwargs):
  68     return BasicEnv().cmd(args, **kwargs)
  69
  70
  71 class VerboseWrapper(object):
  72
  73     def __init__(self, env):
  74         self._env = env
  75
  76     def cmd(self, args, **kwargs):
  77         pprint.pprint(args)
  78         return self._env.cmd(args, **kwargs)
  79
  80
  81 class PrefixCmdEnv(object):
  82
  83     def __init__(self, prefix_cmd, env):
  84         self._prefix_cmd = prefix_cmd
  85         self._env = env
  86
  87     def cmd(self, args, **kwargs):
  88         return self._env.cmd(self._prefix_cmd + args, **kwargs)
  89
  90
  91 def in_dir(dir_path):
  92     return ["sh", "-c", 'cd "$1" && shift && exec "$@"',
  93             "inline_chdir_script", dir_path]
  94
  95
  96 def write_file_cmd(filename, data):
  97     return ["sh", "-c", 'echo -n "$1" >"$2"', "inline_script", data, filename]
  98
  99
 100 def append_file_cmd(filename, data):
 101     return ["sh", "-c", 'echo -n "$1" >>"$2"', "inline_script", data, filename]
 102
 103
 104 def clean_environ_except_home_env(env):
 105     # Resets everything except HOME, so make sure wrapped env sets
 106     # HOME correctly (e.g. sudo -H).
 107     path = "PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin"
 108     return PrefixCmdEnv(
 109         ["sh", "-c", 'env -i HOME="$HOME" %s "$@"' % path,
 110          "clean_environ_env"], env)
 111
 112
 113 def set_environ_vars_env(vars, env):
 114     return PrefixCmdEnv(["env"] + ["%s=%s" % (key, value)
 115                                    for key, value in vars], env)
 116
 117
 118 def get_all_mounts():
 119     fh = open("/proc/mounts", "r")
 120     for line in fh:
 121         parts = line.split()
 122         yield parts[1]
 123     fh.close()
 124
 125
 126 def is_mounted(path):
 127     return os.path.realpath(path) in list(get_all_mounts())
 128
 129
 130 def is_path_below(parent, child):
 131     parent = parent.rstrip("/")
 132     child = child.rstrip("/")
 133     return parent == child or child.startswith(parent + "/")
 134
 135
 136 DELETED_SUFFIX = "\\040(deleted)"
 137
 138
 139 def get_mounts_below(pathname):
 140     prefix = os.path.realpath(pathname)
 141     for mount_path in get_all_mounts():
 142         # The 2.6.22-14 kernel on our buildbot machine seems to add
 143         # this to lines in /proc/mounts if the source directory of the
 144         # bind mount gets deleted.
 145         if mount_path.endswith(DELETED_SUFFIX):
 146             mount_path = mount_path[:-len(DELETED_SUFFIX)]
 147         if is_path_below(prefix, mount_path):
 148             yield mount_path
 149
 150
 151 def unmount_below(as_root, dir_path):
 152     # When you use "mount --rbind", you can end up with bind mounts
 153     # under bind mounts, such as /dev and /dev/pts.  The latter needs
 154     # to get unmounted first, so sort longest first.
 155     mounts = sorted(get_mounts_below(dir_path), key=len, reverse=True)
 156     for mount_path in mounts:
 157         as_root.cmd(["umount", mount_path])
 158
 159
 160 def mount_proc(root_env, chroot_dir):
 161     proc_path = os.path.join(chroot_dir, "proc")
 162     if not is_mounted(proc_path):
 163         root_env.cmd(["mount", "-t", "proc", "proc", proc_path])
 164
 165
 166 def bind_mount_dev_log(root_env, chroot_dir):
 167     # Needed for logging to work inside a chroot
 168     devlog_path = os.path.join(chroot_dir, "dev", "log")
 169     if not is_mounted(devlog_path):
 170         root_env.cmd(["touch", devlog_path])
 171         root_env.cmd(["mount", "--bind", "/dev/log", devlog_path])
 172
 173
 174 def mount_dev_pts(root_env, chroot_dir):
 175     # /dev/ptmx and /dev/pts are needed for openpty() to work.
 176     devpts_path = os.path.join(chroot_dir, "dev", "pts")
 177     if not is_mounted(devpts_path):
 178         root_env.cmd(["mount", "-t", "devpts", "devpts", devpts_path])
 179
 180
 181 def mount_sys(root_env, chroot_dir):
 182     # Needed for activitymonitor in chroots on machines without a static IP
 183     sys_path = os.path.join(chroot_dir, "sys")
 184     if not is_mounted(sys_path):
 185         root_env.cmd(["mount", "-t", "sysfs", "sys", sys_path])
 186
 187
 188 def bind_mount_x11(root_env, chroot_dir):
 189     dest_path = os.path.join(chroot_dir, "tmp/.X11-unix")
 190     if not is_mounted(dest_path):
 191         root_env.cmd(["mount", "--bind", "/tmp/.X11-unix", dest_path])
 192
 193
 194 class EnvWithHook(object):
 195
 196     # Wrapper for an environment with a hook function that gets run
 197     # the first time the wrapper is used.
 198
 199     def __init__(self, hook_func, env):
 200         self._hook_func = hook_func
 201         self._env = env
 202
 203     def cmd(self, args, **kwargs):
 204         self._hook_func()
 205         self._hook_func = lambda: None
 206         return self._env.cmd(args, **kwargs)
 207
 208
 209 def chroot_env(chroot_dir, as_root, environ_vars,
 210                do_forward_x11=False, do_forward_ssh_agent=False):
 211     def hook():
 212         # This reads the local /proc/mounts so will not work properly
 213         # if as_root is remote.
 214         mount_proc(as_root, chroot_dir)
 215         mount_sys(as_root, chroot_dir)
 216         bind_mount_dev_log(as_root, chroot_dir)
 217         mount_dev_pts(as_root, chroot_dir)
 218         if do_forward_x11:
 219             bind_mount_x11(as_root, chroot_dir)
 220             as_root.cmd(["bash", "-c",
 221                          """\
 222 xauth nlist | HOME=/root env -u XAUTHORITY chroot "$1" xauth nmerge -""",
 223                          "inline_chroot_script", chroot_dir])
 224         if do_forward_ssh_agent:
 225             forward_ssh_agent(as_root, environ_vars, chroot_dir)
 226     after_hook = EnvWithHook(hook, as_root)
 227     if do_forward_x11:
 228         # Unset XAUTHORITY so that it is treated as defaulting to
 229         # $HOME/.Xauthority.
 230         return PrefixCmdEnv(["env", "-u", "XAUTHORITY", "HOME=/root",
 231                              "chroot", chroot_dir], after_hook)
 232     else:
 233         return PrefixCmdEnv(["chroot", chroot_dir], after_hook)
 234
 235
 236 # "user" is a string to pass to sudo, or None for root.
 237 def chroot_and_sudo_env(chroot_dir, as_root, environ_vars,
 238                         user, do_forward_x11, do_forward_ssh_agent):
 239     in_chroot = chroot_env(chroot_dir, as_root, environ_vars,
 240                            do_forward_x11=do_forward_x11,
 241                            do_forward_ssh_agent=do_forward_ssh_agent)
 242     if user is not None:
 243         if do_forward_x11:
 244             in_chroot = xsudo_env(user, in_chroot)
 245         else:
 246             in_chroot = PrefixCmdEnv(["sudo", "-H", "-u", user], in_chroot)
 247     in_chroot = clean_environ_except_home_env(in_chroot)
 248     vars_to_keep = {}
 249     if do_forward_x11 and "DISPLAY" in environ_vars:
 250         vars_to_keep["DISPLAY"] = environ_vars["DISPLAY"]
 251     if do_forward_ssh_agent and "SSH_AUTH_SOCK" in environ_vars:
 252         vars_to_keep["SSH_AUTH_SOCK"] = environ_vars["SSH_AUTH_SOCK"]
 253     return set_environ_vars_env(vars_to_keep.items(), in_chroot)
 254
 255
 256 def xsudo_env(user_name, env):
 257     # This copies all Xauthority entries across (using nlist) instead
 258     # of just the one for DISPLAY (which would use nextract).
 259     # "xauth nextract" has a bug when used with TCP displays of the
 260     # form "localhost:N".
 261     return PrefixCmdEnv(["bash", "-c", """
 262 set -e
 263 function my_sudo {
 264   sudo -H -p "Password to get from %u to %U on %H: " -u '"""+user_name+"""' "$@"
 265 }
 266 xauth nlist | my_sudo xauth nmerge -
 267 my_sudo "$@"
 268 """, "inline_sudo_script"], env)
 269
 270
 271 def bash_login_env(env):
 272     # This is useful when running a command directly without starting
 273     # up an interactive shell, assuming that ~/.bash_profile or
 274     # ~/.bashrc are responsible for setting PYTHONPATH etc.
 275     #
 276     # This is loosely equivalent to "bash --login", which makes bash
 277     # load ~/.bash_profile.  This typically sets PATH to include
 278     # conductor/bin.  On typical ThirdPhase stations, ~/.bash_profile
 279     # just sources /etc/conductor/profile.
 280     #
 281     # The default Ubuntu bashrc exits early if PS1 is not set; we set
 282     # PS1 as a workaround.  (bash sets PS1 when started in interactive
 283     # mode; it unsets PS1 when started in non-interactive mode.)
 284     return PrefixCmdEnv(
 285         ["bash", "-c", """
 286 PS1="$ "
 287 if [ -f ~/.bash_profile ]
 288   then
 289     source ~/.bash_profile
 290   else
 291     source ~/.profile
 292 fi
 293 exec "$@"
 294 """,
 295          "inline_shell_script"], env)
 296
 297
 298 def setup_sudo(as_root, username):
 299     # Warning: this overwrites your sudo config
 300     as_root.cmd(write_file_cmd("/etc/sudoers", """\
 301 root ALL=(ALL) ALL
 302 %s ALL=(ALL) NOPASSWD:ALL
 303 """ % username))
 304
 305
 306 def make_relative_to_root(path):
 307     assert path.startswith("/")
 308     return path.lstrip("/")
 309
 310
 311 def bind_mount(as_root, chroot_dir, abs_path):
 312     path = make_relative_to_root(abs_path)
 313     dest_path = os.path.join(chroot_dir, path)
 314     if not is_mounted(dest_path):
 315         as_root.cmd(["mkdir", "-p", dest_path])
 316         as_root.cmd(["mount", "--bind", os.path.join("/", path), dest_path])
 317
 318
 319 def forward_ssh_agent(as_root, environ, chroot_dir):
 320     path = os.path.dirname(environ["SSH_AUTH_SOCK"])
 321     bind_mount(as_root, chroot_dir, path)
 322
 323
 324 def disable_daemons(as_root):
 325     as_root.cmd(write_file_cmd("/usr/sbin/policy-rc.d", "#!/bin/sh\nexit 101"))
 326     as_root.cmd(["chmod", "+x", "/usr/sbin/policy-rc.d"])
 327
 328
 329 def reenable_daemons(as_root):
 330     as_root.cmd(["rm", "/usr/sbin/policy-rc.d"])
 331
 332
 333 class Process(object):
 334
 335     def __init__(self, pid):
 336         self._pid = pid
 337         try:
 338             args = read_file(
 339                 os.path.join("/proc/%i/cmdline" % pid)).split("\0")
 340         except (OSError, IOError), e:
 341             if e.errno == errno.ENOENT:
 342                 # The process may have exited
 343                 args = None
 344             else:
 345                 raise
 346         if args is not None and args[-1] == "":
 347             # There is usually a useless trailing \0 in the cmdline
 348             # file, with the exception of /sbin/init and some programs
 349             # that modify their argv such as avahi-daemon.
 350             args.pop()
 351         self._cmdline = args
 352         try:
 353             self._root_dir = os.readlink(os.path.join("/proc/%i/root" % pid))
 354         except (OSError, IOError), e:
 355             if e.errno == errno.ENOENT:
 356                 # * If the process terminates, /proc/$pid will not exist
 357                 # * readlink can fail with ENOENT even if /proc/$pid/root shows
 358                 #   up in listdir
 359                 self._root_dir = None
 360             elif e.errno == errno.EACCES:
 361                 # If we don't own the process we probably will not be
 362                 # allowed to read this
 363                 self._root_dir = None
 364             else:
 365                 raise
 366
 367     def get_pid(self):
 368         return self._pid
 369
 370     def get_cmdline(self):
 371         return self._cmdline
 372
 373     def get_root_dir(self):
 374         return self._root_dir
 375
 376     def kill(self, signal_number):
 377         """Send the specified signal to the specified process
 378
 379         Returns True if the signal was sent succesfully and False if the
 380         specified process did not exist at the time of the attempt.
 381         """
 382         try:
 383             os.kill(self._pid, signal_number)
 384         except OSError, e:
 385             if e.errno == errno.ESRCH:
 386                 return False
 387             else:
 388                 raise
 389         else:
 390             return True
 391
 392
 393 def list_processes():
 394     for pid_string in os.listdir("/proc"):
 395         try:
 396             pid = int(pid_string)
 397         except ValueError:
 398             pass
 399         else:
 400             yield Process(pid)
 401
 402
 403 def find_chrooted(chroot):
 404     chroot_realpath = os.path.realpath(chroot)
 405     for proc in list_processes():
 406         if (proc.get_root_dir() is not None
 407             and is_path_below(chroot_realpath, proc.get_root_dir())):
 408             yield proc
 409
 410
 411 def kill_chrooted(chroot):
 412     found = set()
 413     while True:
 414         procs = list(find_chrooted(chroot))
 415         if len(procs) == 0:
 416             break
 417         for proc in procs:
 418             if proc.get_pid() not in found:
 419                 print "killing process running in chroot:", \
 420                       proc.get_pid(), proc.get_cmdline()
 421             found.add(proc.get_pid())
 422             proc.kill(signal.SIGKILL)
 423         time.sleep(0.5)