Mini Shell

"""Entry point for /opt/imh-python/bin/python3 -m cproc"""
import sys
import argparse
import signal
from cproc import Proc, ProcLimit


def _sig_arg(val):
    try:
        return signal.Signals[val.upper()].value
    except KeyError as exc:
        raise argparse.ArgumentTypeError("invalid signal") from exc


def main():
    """Entry point for /opt/imh-python/bin/python3 -m cproc"""
    parser = argparse.ArgumentParser(
        description='runs a program while monitoring server load',
        usage='usage: cproc [-h] -l MAX_LOAD -- cmd args here',
    )
    parser.add_argument('-l', dest='max_load', type=float, help='max load')
    parser.add_argument(
        '-m', dest='max_mem', type=int, help='Max memory usage in MiB'
    )
    parser.add_argument(
        '--signal',
        type=_sig_arg,
        default=15,
        help='kill signal if over max memory (default: SIGTERM)',
    )
    parser.add_argument('cmd', nargs='+')
    args = parser.parse_args()
    lim = ProcLimit(args.max_load, max_mem=args.max_mem, mem_signal=args.signal)
    try:
        code = Proc.run(
            args.cmd,
            lim=lim,
            shell=False,
            encoding=None,
            stdin=sys.stdin,
            stdout=sys.stdout,
            stderr=sys.stderr,
            check=False,
        ).returncode
    except KeyboardInterrupt:
        sys.exit('KeyboardInterrupt')
    sys.exit(code)


main()