nepi/util/execfuncs.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
  19 from nepi.util.sshfuncs import ProcStatus, STDOUT, log, shell_escape
  20
  21 import logging
  22 import shlex
  23 import subprocess
  24
  25 from six import PY2
  26
  27 def lexec(command,
  28           user = None,
  29           sudo = False,
  30           env = None):
  31     """
  32     Executes a local command, returns ((stdout, stderr), process)
  33     """
  34     if env:
  35         export = ''
  36         for envkey, envval in env.items():
  37             export += "{}={} ".format(envkey, envval)
  38         command = "{} {}".format(export, command)
  39
  40     if sudo:
  41         command = "sudo {}".format(command)
  42
  43     # XXX: Doing 'su user' blocks waiting for a password on stdin
  44     #elif user:
  45     #    command = "su {} ; {} ".format(user, command)
  46
  47     extras = {} if PY2 else {'universal_newlines' : True}
  48     proc = subprocess.Popen(command,
  49                             shell = True,
  50                             stdout = subprocess.PIPE,
  51                             stderr = subprocess.PIPE,
  52                             **extras)
  53
  54     out = err = ""
  55     log_msg = "lexec - command {} ".format(command)
  56
  57     try:
  58         out, err = proc.communicate()
  59         log(log_msg, logging.DEBUG, out, err)
  60     except:
  61         log(log_msg, logging.ERROR, out, err)
  62         raise
  63
  64     return ((out, err), proc)
  65
  66 def lcopy(source, dest, recursive = False):
  67     """
  68     Copies from/to locally.
  69     """
  70
  71     args = ["cp"]
  72     if recursive:
  73         args.append("-r")
  74
  75     if isinstance(source, list):
  76         args.extend(source)
  77     else:
  78         args.append(source)
  79
  80     if isinstance(dest, list):
  81         args.extend(dest)
  82     else:
  83         args.append(dest)
  84
  85     extras = {} if PY2 else {'universal_newlines' : True}
  86     proc = subprocess.Popen(args,
  87                             stdout=subprocess.PIPE,
  88                             stderr=subprocess.PIPE,
  89                             **extras)
  90
  91     out = err = ""
  92     command = " ".join(args)
  93     log_msg = " lcopy - command {} ".format(command)
  94
  95     try:
  96         out, err = proc.communicate()
  97         log(log_msg, logging.DEBUG, out, err)
  98     except:
  99         log(log_msg, logging.ERROR, out, err)
 100         raise
 101
 102     return ((out, err), proc)
 103
 104 def lspawn(command, pidfile,
 105            stdout = '/dev/null',
 106            stderr = STDOUT,
 107            stdin = '/dev/null',
 108            home = None,
 109            create_home = False,
 110            sudo = False,
 111            user = None):
 112     """
 113     Spawn a local command such that it will continue working asynchronously.
 114
 115     Parameters:
 116         command: the command to run - it should be a single line.
 117
 118         pidfile: path of a (ideally unique to this task) pidfile for tracking the process.
 119
 120         stdout: path of a file to redirect standard output to - must be a string.
 121             Defaults to /dev/null
 122         stderr: path of a file to redirect standard error to - string or the special STDOUT value
 123             to redirect to the same file stdout was redirected to. Defaults to STDOUT.
 124         stdin: path of a file with input to be piped into the command's standard input
 125
 126         home: path of a folder to use as working directory - should exist, unless you specify create_home
 127
 128         create_home: if True, the home folder will be created first with mkdir -p
 129
 130         sudo: whether the command needs to be executed as root
 131
 132     Returns:
 133         (stdout, stderr), process
 134
 135         Of the spawning process, which only captures errors at spawning time.
 136         Usually only useful for diagnostics.
 137     """
 138     # Start process in a "daemonized" way, using nohup and heavy
 139     # stdin/out redirection to avoid connection issues
 140     if stderr is STDOUT:
 141         stderr = '&1'
 142     else:
 143         stderr = ' ' + stderr
 144
 145     escaped_pidfile = shell_escape(pidfile)
 146     daemon_command = '{{ {{ {command}  > {stdout} 2>{stderr} < {stdin} & }} ; echo $! 1 > {escaped_pidfile} ; }}'\
 147                      .format(**locals())
 148
 149     cmd = "{create}{gohome} rm -f {pidfile} ; {sudo} bash -c {command} "\
 150           .format(command = shell_escape(daemon_command),
 151                   sudo = 'sudo -S' if sudo else '',
 152                   pidfile = shell_escape(pidfile),
 153                   gohome = 'cd {} ; '.format(shell_escape(home)) if home else '',
 154                   create = 'mkdir -p {} ; '.format(shell_escape(home)) if create_home else '')
 155
 156     (out, err), proc = lexec(cmd)
 157
 158     if proc.wait():
 159         raise RuntimeError("Failed to set up local application: {} {}"
 160                            .format(out, err))
 161
 162     return ((out, err), proc)
 163
 164 def lgetpid(pidfile):
 165     """
 166     Check the pidfile of a process spawned with remote_spawn.
 167
 168     Parameters:
 169         pidfile: the pidfile passed to remote_span
 170
 171     Returns:
 172
 173         A (pid, ppid) tuple useful for calling remote_status and remote_kill,
 174         or None if the pidfile isn't valid yet (maybe the process is still starting).
 175     """
 176
 177     (out, err), proc = lexec("cat {}".format(pidfile))
 178
 179     if proc.wait():
 180         return None
 181
 182     if out:
 183         try:
 184             return [ int(x) for x in out.strip().split(' ', 1) ]
 185         except:
 186             # Ignore, many ways to fail that don't matter that much
 187             return None
 188
 189 def lstatus(pid, ppid):
 190     """
 191     Check the status of a process spawned with remote_spawn.
 192
 193     Parameters:
 194         pid/ppid: pid and parent-pid of the spawned process. See remote_check_pid
 195
 196     Returns:
 197
 198         One of NOT_STARTED, RUNNING, FINISHED
 199     """
 200
 201     (out, err), proc = lexec(
 202         # Check only by pid. pid+ppid does not always work (especially with sudo)
 203         " (( ps --pid {pid} -o pid | grep -c {pid} && echo 'wait')  || echo 'done' ) | tail -n 1"
 204         .format(ppid = ppid, pid = pid))
 205
 206     if proc.wait():
 207         return ProcStatus.NOT_STARTED
 208
 209     status = False
 210     if out:
 211         status = (out.strip() == 'wait')
 212     else:
 213         return ProcStatus.NOT_STARTED
 214
 215     return ProcStatus.RUNNING if status else ProcStatus.FINISHED
 216
 217 def lkill(pid, ppid, sudo = False, nowait = False):
 218     """
 219     Kill a process spawned with lspawn.
 220
 221     First tries a SIGTERM, and if the process does not end in 10 seconds,
 222     it sends a SIGKILL.
 223
 224     Parameters:
 225         pid/ppid: pid and parent-pid of the spawned process. See remote_check_pid
 226
 227         sudo: whether the command was run with sudo - careful killing like this.
 228
 229     Returns:
 230
 231         Nothing, should have killed the process
 232     """
 233
 234     subkill = "$(ps --ppid {pid} -o pid h)".format(pid=pid)
 235     cmd_format = """
 236 SUBKILL="{subkill}" ;
 237 {sudo} kill -- -{pid} $SUBKILL || /bin/true
 238 {sudo} kill {pid} $SUBKILL || /bin/true
 239 for x in 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 ; do
 240     sleep 0.2
 241     if [ `ps --pid {pid} -o pid | grep -c {pid}` == '0' ]; then
 242         break
 243     else
 244         {sudo} kill -- -{pid} $SUBKILL || /bin/true
 245         {sudo} kill {pid} $SUBKILL || /bin/true
 246     fi
 247     sleep 1.8
 248 done
 249 if [ `ps --pid {pid} -o pid | grep -c {pid}` != '0' ]; then
 250     {sudo} kill -9 -- -{pid} $SUBKILL || /bin/true
 251     {sudo} kill -9 {pid} $SUBKILL || /bin/true
 252 fi
 253 """
 254     if nowait:
 255         cmd = "( {} ) >/dev/null 2>/dev/null </dev/null &".format(cmd)
 256
 257     (out, err), proc = lexec(
 258         cmd_format.format(
 259             ppid = ppid,
 260             pid = pid,
 261             sudo = 'sudo -S' if sudo else '',
 262             subkill = subkill))
 263
 264