#!/usr/bin/env python3

import argparse
import resource
import sys
import os
import subprocess


def main():
    parser = argparse.ArgumentParser(description="Enable or disable coredump for the PIDs")
    parser.add_argument("-e", "--enable", action="store_true")
    parser.add_argument("-c", "--command", action="store_true", help="exec mode")
    parser.add_argument(
        metavar='PID/program/arg', dest='free_args', nargs="*",
        help="command line to run in exec mode, PID otherwise. If no PID is given, will read from stdin")
    args = parser.parse_args()

    limit = (
        (0, resource.RLIM_INFINITY)[args.enable],
        resource.RLIM_INFINITY
    )

    if args.command:
        resource.prlimit(os.getpid(), resource.RLIMIT_CORE, limit)
        ret = subprocess.Popen(args.free_args).wait()
        if ret != 0:
            raise Exception(f"{args.free_args} exited with code {ret}")
    else:
        if not args.free_args:
            for pid in sys.stdin:
                args.free_args.append(int(pid.strip()))
        for pid in args.free_args:
            resource.prlimit(int(pid), resource.RLIMIT_CORE, limit)


if __name__ == "__main__":
    main()
