ColdCore/daemon_manager.py

50 lines
1.7 KiB
Python

#!/usr/bin/python3
# from picoctf platform
import glob, imp, argparse, time
from os.path import splitext, basename
import config
def load_modules(directory):
files = glob.glob("{}/*.py".format(directory))
return [imp.load_source(splitext(basename(module))[0], module) for module in files]
def run_modules(modules, interval):
while True:
start_time = time.time()
for module in modules:
module.run()
time.sleep(max(interval - (time.time() - start_time), 0))
def main():
parser = argparse.ArgumentParser(description="{} daemon manager".format(config.ctf_name))
parser.add_argument("-l", action="store_true", dest="show_list", help="List all daemons")
parser.add_argument("-a", "--all", dest="run_all", action="store_true", help="Run all daemons")
parser.add_argument("-i", "--interval", action="store", type=int, help="The interval in which to run the daemons", default=60)
parser.add_argument("-d", "--daemon-directory", action="store", help="The directory which contains the daemons", default="daemons")
parser.add_argument("modules", nargs="*", help="The daemon modules to run")
args = parser.parse_args()
modules = load_modules(args.daemon_directory)
if args.show_list:
for module in modules:
print(module.__name__)
elif args.run_all:
run_modules(modules, args.interval)
else:
if len(args.modules) == 0:
parser.print_help()
exit(1)
selected_modules = [m for m in modules if m.__name__ in args.modules]
if len(selected_modules) == 0:
parser.print_help()
exit(1)
run_modules(selected_modules, args.interval)
main()