shell.py 4.41 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
has_resource_module = True
try:
Johannes Bechberger's avatar
Johannes Bechberger committed
12
    import resource
13
except ImportError:
Johannes Bechberger's avatar
Johannes Bechberger committed
14
15
    has_resource_module = False
    pass
16
	
Johannes Bechberger's avatar
Johannes Bechberger committed
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
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):
35
36
    if not has_resource_module:
        return
Johannes Bechberger's avatar
Johannes Bechberger committed
37
38
39
40
41
42
43
44
45
    (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):
46
    def __init__(self, cmd, timeout, env, rlimit, input_str):
Johannes Bechberger's avatar
Johannes Bechberger committed
47
48
49
50
51
52
        self.cmd       = cmd
        self.timeout   = timeout
        self.env       = env
        self.proc      = None
        self.exception = None
        self.rlimit    = rlimit
53
        self.input_str = input_str
Johannes Bechberger's avatar
Johannes Bechberger committed
54
55
56
57
58
59
60
61
62
63
64
        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):
65
66
        if not has_resource_module:
            return
Johannes Bechberger's avatar
Johannes Bechberger committed
67
68
69
70
71
72
73
        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:
74
75
76
            prexec_args = {}
            if has_resource_module:
                prexec_args["preexec_fn"] = self._set_rlimit
Johannes Bechberger's avatar
Johannes Bechberger committed
77
78
79
            self.proc = subprocess.Popen(self.cmd,
                                         stdout=subprocess.PIPE,
                                         stderr=subprocess.PIPE,
80
                                         stdin=subprocess.PIPE if self.input_str else None,
81
                                         env=self.env,
Johannes Bechberger's avatar
Johannes Bechberger committed
82
                                         shell=True,
83
                                         **prexec_args)
84
85
86
87
            #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
88
89
            self.out, self.err = x
            self.returncode = self.proc.returncode
90
91
92
93
94
95
        except subprocess.TimeoutExpired as t:
            self.returncode = -signal.SIGXCPU
            self.out = bytearray()
            self.err = bytearray()
        #except Exception as e:
        #    self.exception = e
Johannes Bechberger's avatar
Johannes Bechberger committed
96
97

    def run(self):
98
        self._run_process()
Johannes Bechberger's avatar
Johannes Bechberger committed
99
100
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
        if self.returncode in _EXIT_CODES:
            _LOG.debug("returncode %d recognized as '%s'" % (self.returncode, _EXIT_CODES[self.returncode]))
108
            raise SigKill(-self.returncode, _EXIT_CODES[self.returncode] + ": " + os.strerror(-self.returncode))
Johannes Bechberger's avatar
Johannes Bechberger committed
109
110
        return (self.out, self.err, self.returncode)

111
def execute(cmd, env=None, timeout=0, rlimit=None, propagate_sigint=True, input_str = None):
Johannes Bechberger's avatar
Johannes Bechberger committed
112
113
114
115
    """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
116
117
        pass
        #cmd = filter(lambda x: x, cmd.split(' '))
118
    else:
Johannes Bechberger's avatar
Johannes Bechberger committed
119
        #cmd = shlex.split(cmd[0]) + cmd[1:]
Johannes Bechberger's avatar
Fix #36    
Johannes Bechberger committed
120
        cmd = " ".join(shlex.quote(c) for c in cmd)
121
    exc = _Execute(cmd, timeout, env, rlimit, input_str)
Johannes Bechberger's avatar
Johannes Bechberger committed
122
123
124
125
126
127
128
129
130
    (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)