shell.py 4.45 KB
Newer Older
Johannes Bechberger's avatar
Johannes Bechberger committed
1
# source: https://github.com/libfirm/sisyphus, but slightly modified to support Windows
Johannes Bechberger's avatar
Johannes Bechberger committed
2 3 4 5 6

"""
Convenience function
Alternative to subprocess and os.system
"""
7
import os
Johannes Bechberger's avatar
Johannes Bechberger committed
8
import shlex
Johannes Bechberger's avatar
Johannes Bechberger committed
9
import subprocess
10 11 12

import errno

13 14
has_resource_module = True
try:
Johannes Bechberger's avatar
Johannes Bechberger committed
15
    import resource
16
except ImportError:
Johannes Bechberger's avatar
Johannes Bechberger committed
17 18
    has_resource_module = False
    pass
19
	
Johannes Bechberger's avatar
Johannes Bechberger committed
20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37
import sys
import signal
import threading

import logging
_LOG = logging.getLogger("sisyphus")

_EXIT_CODES = dict((-k, v) for v, k in signal.__dict__.items() if v.startswith('SIG'))
del _EXIT_CODES[0]


class SigKill(Exception):
    def __init__(self, retcode, name):
        self.retcode = retcode
        self.name = name


def _lower_rlimit(res, limit):
38 39
    if not has_resource_module:
        return
Johannes Bechberger's avatar
Johannes Bechberger committed
40 41 42 43 44 45 46 47 48
    (soft, hard) = resource.getrlimit(res)
    if soft > limit or soft == resource.RLIM_INFINITY:
        soft = limit
    if hard > limit or hard == resource.RLIM_INFINITY:
        hard = limit
    resource.setrlimit(res, (soft, hard))


class _Execute(object):
49
    def __init__(self, cmd, timeout, env, rlimit, input_str):
Johannes Bechberger's avatar
Johannes Bechberger committed
50 51 52 53 54 55
        self.cmd       = cmd
        self.timeout   = timeout
        self.env       = env
        self.proc      = None
        self.exception = None
        self.rlimit    = rlimit
56
        self.input_str = input_str
Johannes Bechberger's avatar
Johannes Bechberger committed
57 58 59 60 61 62 63 64 65 66 67
        MB = 1024 * 1024
        if not 'RLIMIT_CORE' in rlimit:
            rlimit['RLIMIT_CORE'] = 0
        if not 'RLIMIT_DATA' in rlimit:
            rlimit['RLIMIT_DATA'] = 1024 * MB
        if not 'RLIMIT_STACK' in rlimit:
            rlimit['RLIMIT_STACK'] = 1024 * MB
        if not 'RLIMIT_FSIZE' in rlimit:
            rlimit['RLIMIT_FSIZE'] = 512 * MB

    def _set_rlimit(self):
68 69
        if not has_resource_module:
            return
Johannes Bechberger's avatar
Johannes Bechberger committed
70 71 72 73 74 75 76
        if self.timeout > 0.0:
            _lower_rlimit(resource.RLIMIT_CPU, self.timeout)
        for k,v in self.rlimit.items():
            _lower_rlimit(getattr(resource,k), v)

    def _run_process(self):
        try:
77 78 79
            prexec_args = {}
            if has_resource_module:
                prexec_args["preexec_fn"] = self._set_rlimit
Johannes Bechberger's avatar
Johannes Bechberger committed
80 81 82
            self.proc = subprocess.Popen(self.cmd,
                                         stdout=subprocess.PIPE,
                                         stderr=subprocess.PIPE,
83
                                         stdin=subprocess.PIPE if self.input_str else None,
84
                                         env=self.env,
Johannes Bechberger's avatar
Johannes Bechberger committed
85
                                         shell=True,
86
                                         **prexec_args)
87 88 89 90
            #if self.input_str:
            #    self.proc.stdin.write(self.input_str.decode)
            input_bytes = self.input_str.encode() if self.input_str else None
            x = self.proc.communicate(input=input_bytes, timeout=self.timeout if self.timeout > 0.0 else None)
Johannes Bechberger's avatar
Johannes Bechberger committed
91 92
            self.out, self.err = x
            self.returncode = self.proc.returncode
93
        except subprocess.TimeoutExpired as t:
94 95
            self.returncode = -errno.ETIMEDOUT
            raise SigKill(signal.SIGXCPU, "SIGXCPU")
96 97
        #except Exception as e:
        #    self.exception = e
Johannes Bechberger's avatar
Johannes Bechberger committed
98 99

    def run(self):
100
        self._run_process()
Johannes Bechberger's avatar
Johannes Bechberger committed
101 102 103 104 105 106 107

        # Usually python can recognize application terminations triggered by
        # signals, but somehow it doesn't do this for java (I suspect, that java
        # catches the signals itself but then shuts down more 'cleanly' than it
        # should. Calculate to python convention returncode
        if self.returncode > 127:
            self.returncode = 128 - self.returncode
108
        if self.returncode in _EXIT_CODES or self.returncode == -errno.ETIMEDOUT:
Johannes Bechberger's avatar
Johannes Bechberger committed
109
            _LOG.debug("returncode %d recognized as '%s'" % (self.returncode, _EXIT_CODES[self.returncode]))
110
            raise SigKill(-self.returncode, _EXIT_CODES[self.returncode] + ": " + os.strerror(-self.returncode))
Johannes Bechberger's avatar
Johannes Bechberger committed
111 112
        return (self.out, self.err, self.returncode)

113
def execute(cmd, env=None, timeout=0, rlimit=None, propagate_sigint=True, input_str = None):
Johannes Bechberger's avatar
Johannes Bechberger committed
114 115 116 117
    """Execute a command and return stderr and stdout data"""
    if not rlimit:
        rlimit = dict()
    if cmd is str:
Johannes Bechberger's avatar
Johannes Bechberger committed
118 119
        pass
        #cmd = filter(lambda x: x, cmd.split(' '))
120
    else:
Johannes Bechberger's avatar
Johannes Bechberger committed
121
        #cmd = shlex.split(cmd[0]) + cmd[1:]
Johannes Bechberger's avatar
Fix #36  
Johannes Bechberger committed
122
        cmd = " ".join(shlex.quote(c) for c in cmd)
123
    exc = _Execute(cmd, timeout, env, rlimit, input_str)
Johannes Bechberger's avatar
Johannes Bechberger committed
124 125 126 127 128 129 130 131 132
    (out, err, returncode) = exc.run()
    if returncode == -signal.SIGINT:
        raise KeyboardInterrupt
    return (out, err, returncode)


if __name__ == "__main__":
    out, err, retcode = execute("hostname")
    print (out)