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 += '%s=%s ' % (envkey, envval)
  38         command = "%s %s" % (export, command)
  39
  40     if sudo:
  41         command = "sudo %s" % command
  42
  43     # XXX: Doing 'su user' blocks waiting for a password on stdin
  44     #elif user:
  45     #    command = "su %s ; %s " % (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 %s " % 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 localy.
  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 %s " % 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     daemon_command = '{ { %(command)s  > %(stdout)s 2>%(stderr)s < %(stdin)s & } ; echo $! 1 > %(pidfile)s ; }' % {
 146         'command' : command,
 147         'pidfile' : shell_escape(pidfile),
 148         'stdout' : stdout,
 149         'stderr' : stderr,
 150         'stdin' : stdin,
 151     }
 152
 153     cmd = "%(create)s%(gohome)s rm -f %(pidfile)s ; %(sudo)s bash -c %(command)s " % {
 154             'command' : shell_escape(daemon_command),
 155             'sudo' : 'sudo -S' if sudo else '',
 156             'pidfile' : shell_escape(pidfile),
 157             'gohome' : 'cd %s ; ' % (shell_escape(home),) if home else '',
 158             'create' : 'mkdir -p %s ; ' % (shell_escape(home),) if create_home else '',
 159         }
 160
 161     (out,err), proc = lexec(cmd)
 162
 163     if proc.wait():
 164         raise RuntimeError("Failed to set up application on host %s: %s %s" % (host, out,err,))
 165
 166     return ((out,err), proc)
 167
 168 def lgetpid(pidfile):
 169     """
 170     Check the pidfile of a process spawned with remote_spawn.
 171
 172     Parameters:
 173         pidfile: the pidfile passed to remote_span
 174
 175     Returns:
 176
 177         A (pid, ppid) tuple useful for calling remote_status and remote_kill,
 178         or None if the pidfile isn't valid yet (maybe the process is still starting).
 179     """
 180
 181     (out,err), proc = lexec("cat %s" % pidfile )
 182
 183     if proc.wait():
 184         return None
 185
 186     if out:
 187         try:
 188             return [ int(x) for x in out.strip().split(' ', 1) ]
 189         except:
 190             # Ignore, many ways to fail that don't matter that much
 191             return None
 192
 193 def lstatus(pid, ppid):
 194     """
 195     Check the status of a process spawned with remote_spawn.
 196
 197     Parameters:
 198         pid/ppid: pid and parent-pid of the spawned process. See remote_check_pid
 199
 200     Returns:
 201
 202         One of NOT_STARTED, RUNNING, FINISHED
 203     """
 204
 205     (out,err), proc = lexec(
 206         # Check only by pid. pid+ppid does not always work (especially with sudo)
 207         " (( ps --pid %(pid)d -o pid | grep -c %(pid)d && echo 'wait')  || echo 'done' ) | tail -n 1" % {
 208             'ppid' : ppid,
 209             'pid' : pid,
 210         })
 211
 212     if proc.wait():
 213         return ProcStatus.NOT_STARTED
 214
 215     status = False
 216     if out:
 217         status = (out.strip() == 'wait')
 218     else:
 219         return ProcStatus.NOT_STARTED
 220
 221     return ProcStatus.RUNNING if status else ProcStatus.FINISHED
 222
 223 def lkill(pid, ppid, sudo = False, nowait = False):
 224     """
 225     Kill a process spawned with lspawn.
 226
 227     First tries a SIGTERM, and if the process does not end in 10 seconds,
 228     it sends a SIGKILL.
 229
 230     Parameters:
 231         pid/ppid: pid and parent-pid of the spawned process. See remote_check_pid
 232
 233         sudo: whether the command was run with sudo - careful killing like this.
 234
 235     Returns:
 236
 237         Nothing, should have killed the process
 238     """
 239
 240     subkill = "$(ps --ppid %(pid)d -o pid h)" % { 'pid' : pid }
 241     cmd = """
 242 SUBKILL="%(subkill)s" ;
 243 %(sudo)s kill -- -%(pid)d $SUBKILL || /bin/true
 244 %(sudo)s kill %(pid)d $SUBKILL || /bin/true
 245 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
 246     sleep 0.2
 247     if [ `ps --pid %(pid)d -o pid | grep -c %(pid)d` == '0' ]; then
 248         break
 249     else
 250         %(sudo)s kill -- -%(pid)d $SUBKILL || /bin/true
 251         %(sudo)s kill %(pid)d $SUBKILL || /bin/true
 252     fi
 253     sleep 1.8
 254 done
 255 if [ `ps --pid %(pid)d -o pid | grep -c %(pid)d` != '0' ]; then
 256     %(sudo)s kill -9 -- -%(pid)d $SUBKILL || /bin/true
 257     %(sudo)s kill -9 %(pid)d $SUBKILL || /bin/true
 258 fi
 259 """
 260     if nowait:
 261         cmd = "( %s ) >/dev/null 2>/dev/null </dev/null &" % (cmd,)
 262
 263     (out,err),proc = lexec(
 264         cmd % {
 265             'ppid' : ppid,
 266             'pid' : pid,
 267             'sudo' : 'sudo -S' if sudo else '',
 268             'subkill' : subkill,
 269         })
 270
 271