monkey/chaos_monkey/main.py

117 lines
3.8 KiB
Python

import os
import sys
import logging
import traceback
import logging.config
from config import WormConfiguration, EXTERNAL_CONFIG_FILE
from model import MONKEY_ARG, DROPPER_ARG
from dropper import MonkeyDrops
from monkey import ChaosMonkey
import argparse
import json
if __name__ == "__main__":
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
__author__ = 'itamar'
LOG = None
LOG_CONFIG = {'version': 1,
'disable_existing_loggers': False,
'formatters': {'standard': {'format': '%(asctime)s [%(process)d:%(levelname)s] %(module)s.%(funcName)s.%(lineno)d: %(message)s'},
},
'handlers': {'console': {'class': 'logging.StreamHandler',
'level': 'DEBUG',
'formatter': 'standard'},
'file': {'class': 'logging.FileHandler',
'level': 'DEBUG',
'formatter': 'standard',
'filename': None}
},
'root': {'level': 'DEBUG',
'handlers': ['console']},
}
def main():
global LOG
if 2 > len(sys.argv):
return True
monkey_mode = sys.argv[1]
if not monkey_mode in [MONKEY_ARG, DROPPER_ARG]:
return True
config_file = EXTERNAL_CONFIG_FILE
arg_parser = argparse.ArgumentParser()
arg_parser.add_argument('-c', '--config')
opts, monkey_args = arg_parser.parse_known_args(sys.argv[2:])
if opts.config:
config_file = opts.config
if os.path.isfile(config_file):
# using print because config can also change log locations
print "Loading config from %s." % config_file
try:
with open(config_file) as config_fo:
json_dict = json.load(config_fo)
WormConfiguration.from_dict(json_dict)
except ValueError as e:
print "Error loading config: %s, using default" % (e,)
else:
print("Config file wasn't supplied and default path: %s wasn't found, using internal default" % (config_file,))
print "Loaded Configuration: %r" % WormConfiguration.as_dict()
try:
if MONKEY_ARG == monkey_mode:
log_path = os.path.expandvars(WormConfiguration.monkey_log_path_windows) if sys.platform == "win32" else WormConfiguration.monkey_log_path_linux
monkey_cls = ChaosMonkey
elif DROPPER_ARG == monkey_mode:
log_path = os.path.expandvars(WormConfiguration.dropper_log_path_windows) if sys.platform == "win32" else WormConfiguration.dropper_log_path_linux
monkey_cls = MonkeyDrops
else:
return True
except ValueError:
return True
if WormConfiguration.use_file_logging:
LOG_CONFIG['handlers']['file']['filename'] = log_path
LOG_CONFIG['root']['handlers'].append('file')
else:
del LOG_CONFIG['handlers']['file']
logging.config.dictConfig(LOG_CONFIG)
LOG = logging.getLogger()
def log_uncaught_exceptions(ex_cls, ex, tb):
LOG.critical(''.join(traceback.format_tb(tb)))
LOG.critical('{0}: {1}'.format(ex_cls, ex))
sys.excepthook = log_uncaught_exceptions
LOG.info(">>>>>>>>>> Initializing monkey (%s): PID %s <<<<<<<<<<",
monkey_cls.__name__, os.getpid())
monkey = monkey_cls(monkey_args)
monkey.initialize()
try:
monkey.start()
if WormConfiguration.serialize_config:
with open(config_file, 'w') as config_fo:
json_dict = WormConfiguration.as_dict()
json.dump(json_dict, config_fo, skipkeys=True, sort_keys=True, indent=4, separators=(',', ': '))
return True
finally:
monkey.cleanup()
if "__main__" == __name__:
if not main():
sys.exit(1)