#!/usr/bin/env python

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 += string.atoi(v)
            if o == '-t':
                self.timelimit = string.atoi(v)
            if o == '-m':
                self.memlimit = string.atoi(v) * 1024
            if o == '-d':
                self.rundir = v
            if o == '-o':
                self.writeto = v
            if o == '-T':
                self.timetime = string.atoi(v)
            if o == '-p':
                self.usepickle = True

        v = os.getenv('GUARD_RLIMIT_OFILE')
        if v: self.ofile = string.atoi(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, e:
                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):
            cmdline = file(procdir + '/cmdline', 'r').readlines()

            # do not get memory usage of this script after just fork
            if len(cmdline) > 0 and \
               string.strip(cmdline[0], '\0') != \
               string.join(self.args[1:], '\0'):
                return

            procstatus = file(procdir + '/status', 'r')
            rss = 0; data = 0; stack = 0
            for line in procstatus:
                n = line[0:6]
                if n == 'VmRSS:':
                    rss = string.atoi(line[7:-3])
                if n == 'VmData':
                    data = string.atoi(line[8:-3])
                if n == 'VmStk:':
                    stack = string.atoi(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:
            f = file(self.writeto, 'w')

        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 >>f, "exitcode: %d" % self.exitcode
            print >>f, "sig: %d" % self.sig
            print >>f, "time: %.3f" % self.timeused
            print >>f, "rss: %d" % self.memrss
            print >>f, "data: %d" % self.memdata
            print >>f, "stack: %d" % self.memstack

        if self.writeto != None: f.close()

if __name__ == '__main__':
    os.umask(0002)
    RunGuard().run()