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