160 lines
5.1 KiB
Python
Executable File
160 lines
5.1 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
|
|
import os, sys, string, signal, resource, time, getopt, pickle
|
|
|
|
class RunGuard:
|
|
|
|
def __init__(self):
|
|
self.timelimit = 1
|
|
self.memlimit = 65536 * 1024
|
|
self.timetime = 5
|
|
self.args = None
|
|
self.writeto = None
|
|
self.nproc = 1
|
|
self.ofile = 32
|
|
self.ldpreload = None
|
|
self.rundir = None
|
|
|
|
self.usepickle = False
|
|
self.memrss = 0
|
|
self.memdata = 0
|
|
self.memstack = 0
|
|
self.timeused = 0
|
|
self.exitcode = 0
|
|
self.sig = 0
|
|
|
|
def run(self):
|
|
self.parse_opts()
|
|
|
|
self.childpid = os.fork()
|
|
if self.childpid == 0:
|
|
self.execute()
|
|
else:
|
|
self.monitor()
|
|
|
|
self.write_result()
|
|
|
|
def parse_opts(self):
|
|
try:
|
|
pos = sys.argv.index('-x')
|
|
optlist, self.args = getopt.gnu_getopt(sys.argv[:pos], 'e:t:m:d:o:T:p')
|
|
self.args = sys.argv[pos:]
|
|
except ValueError:
|
|
optlist, self.args = getopt.gnu_getopt(sys.argv, 'e:t:m:d:o:T:p')
|
|
for o, v in optlist:
|
|
if o == '-e':
|
|
self.nproc += int(v)
|
|
if o == '-t':
|
|
self.timelimit = int(v)
|
|
if o == '-m':
|
|
self.memlimit = int(v) * 1024
|
|
if o == '-d':
|
|
self.rundir = v
|
|
if o == '-o':
|
|
self.writeto = v
|
|
if o == '-T':
|
|
self.timetime = int(v)
|
|
if o == '-p':
|
|
self.usepickle = True
|
|
|
|
v = os.getenv('GUARD_RLIMIT_OFILE')
|
|
if v: self.ofile = int(v)
|
|
self.ldpreload = os.getenv('GUARD_LD_PRELOAD')
|
|
|
|
def execute(self):
|
|
if self.rundir != None and os.path.isdir(self.rundir):
|
|
os.chdir(self.rundir)
|
|
|
|
if self.nproc:
|
|
resource.setrlimit(resource.RLIMIT_NPROC, (self.nproc, self.nproc))
|
|
if self.ofile:
|
|
resource.setrlimit(resource.RLIMIT_OFILE, (self.ofile, self.ofile))
|
|
if self.ldpreload:
|
|
os.putenv('LD_PRELOAD', self.ldpreload)
|
|
|
|
resource.setrlimit(resource.RLIMIT_CPU, (self.timelimit, self.timelimit))
|
|
resource.setrlimit(resource.RLIMIT_AS, (self.memlimit, self.memlimit))
|
|
os.execv(self.args[1], self.args[1:])
|
|
sys.exit(127) # exit with JGE
|
|
|
|
def monitor(self):
|
|
pid = 0
|
|
remaintime = self.timelimit * self.timetime
|
|
while remaintime > 0:
|
|
pid, status, ru = os.wait4(self.childpid, os.WNOHANG)
|
|
self._get_memused()
|
|
if pid > 0: break
|
|
time.sleep(0.05)
|
|
remaintime -= 0.05
|
|
|
|
while pid == 0:
|
|
try:
|
|
os.kill(self.childpid, signal.SIGKILL)
|
|
except OSError:
|
|
pass
|
|
pid, status, ru = os.wait4(self.childpid, os.WNOHANG)
|
|
time.sleep(0.1)
|
|
|
|
if os.WIFEXITED(status):
|
|
self.exitcode = os.WEXITSTATUS(status)
|
|
|
|
if os.WIFSIGNALED(status):
|
|
self.sig = os.WTERMSIG(status)
|
|
|
|
self.timeused = ru[0] + ru[1]
|
|
|
|
def _get_memused(self):
|
|
procdir = '/proc/%d' % self.childpid
|
|
if os.path.isdir(procdir):
|
|
with open(procdir + '/cmdline', 'r') as f:
|
|
cmdline = f.readlines()
|
|
|
|
# do not get memory usage of this script after just fork
|
|
if len(cmdline) > 0 and \
|
|
cmdline[0].strip('\0') != \
|
|
'\0'.join(self.args[1:]):
|
|
return
|
|
|
|
with open(procdir + '/status', 'r') as procstatus:
|
|
rss = 0
|
|
data = 0
|
|
stack = 0
|
|
for line in procstatus:
|
|
n = line[0:6]
|
|
if n == 'VmRSS:':
|
|
rss = int(line[7:-3])
|
|
if n == 'VmData':
|
|
data = int(line[8:-3])
|
|
if n == 'VmStk:':
|
|
stack = int(line[7:-3])
|
|
self.memrss = max(self.memrss, rss)
|
|
if self.memdata + self.memstack < data + stack:
|
|
self.memdata = data
|
|
self.memstack = stack
|
|
return
|
|
|
|
def write_result(self):
|
|
if self.writeto == None:
|
|
f = sys.stdout
|
|
else:
|
|
with open(self.writeto, 'w') as f:
|
|
if self.usepickle:
|
|
obj = { 'exitcode' : self.exitcode,
|
|
'sig' : self.sig,
|
|
'timeused' : self.timeused,
|
|
'memrss' : self.memrss,
|
|
'memdata' : self.memdata,
|
|
'memstack' : self.memstack }
|
|
pickle.dump(obj, f)
|
|
else:
|
|
print("exitcode: %d" % self.exitcode, file=f)
|
|
print("sig: %d" % self.sig, file=f)
|
|
print("time: %.3f" % self.timeused, file=f)
|
|
print("rss: %d" % self.memrss, file=f)
|
|
print("data: %d" % self.memdata, file=f)
|
|
print("stack: %d" % self.memstack, file=f)
|
|
|
|
if __name__ == '__main__':
|
|
os.umask(0o002)
|
|
RunGuard().run()
|