diff --git a/app/api/v2/managers/operation_api_manager.py b/app/api/v2/managers/operation_api_manager.py index fbb37930d..0a20dd8f0 100644 --- a/app/api/v2/managers/operation_api_manager.py +++ b/app/api/v2/managers/operation_api_manager.py @@ -234,20 +234,27 @@ async def get_hosts(self, operation: dict): 'host': tmp_agent.get('host'), 'host_ip_addrs': tmp_agent.get('host_ip_addrs'), 'platform': tmp_agent.get('platform'), - 'reachable_hosts': await self.get_reachable_hosts(operation) + 'reachable_hosts': await self.get_reachable_hosts(agent=tmp_agent) } hosts[host] = tmp_host return hosts - async def get_reachable_hosts(self, operation: dict): + async def get_reachable_hosts(self, agent: dict = None, operation: dict = None): + """ + NOTE: When agent is supplied, only hosts discovered by agent + are retrieved. + """ trait_names = BaseWorld.get_config('reachable_host_traits') or [] paws = () - for agent in operation.get('host_group', []): - paw = agent.get('paw') - if paw: - paws = paws + (paw,) + if agent is not None: + paws = paws + (agent.get('paw'),) + else: + for agent in operation.get('host_group', []): + paw = agent.get('paw') + if paw: + paws = paws + (paw,) hosts = [] for trait in trait_names: diff --git a/server.py b/server.py index 903105fa4..a1b44ecc4 100644 --- a/server.py +++ b/server.py @@ -35,11 +35,13 @@ def setup_logger(level=logging.DEBUG): - logging.basicConfig(level=level, - format='%(asctime)s - %(levelname)-5s (%(filename)s:%(lineno)s %(funcName)s) %(message)s', - datefmt='%Y-%m-%d %H:%M:%S') + logging.basicConfig( + level=level, + format="%(asctime)s - %(levelname)-5s (%(filename)s:%(lineno)s %(funcName)s) %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + ) for logger_name in logging.root.manager.loggerDict.keys(): - if logger_name in ('aiohttp.server', 'asyncio'): + if logger_name in ("aiohttp.server", "asyncio"): continue else: logging.getLogger(logger_name).setLevel(100) @@ -48,10 +50,12 @@ def setup_logger(level=logging.DEBUG): async def start_server(): - await auth_svc.apply(app_svc.application, BaseWorld.get_config('users')) + await auth_svc.apply(app_svc.application, BaseWorld.get_config("users")) runner = web.AppRunner(app_svc.application) await runner.setup() - await web.TCPSite(runner, BaseWorld.get_config('host'), BaseWorld.get_config('port')).start() + await web.TCPSite( + runner, BaseWorld.get_config("host"), BaseWorld.get_config("port") + ).start() def run_tasks(services, run_vue_server=False): @@ -61,8 +65,16 @@ def run_tasks(services, run_vue_server=False): loop.run_until_complete(knowledge_svc.restore_state()) loop.run_until_complete(app_svc.register_contacts()) loop.run_until_complete(app_svc.load_plugins(args.plugins)) - loop.run_until_complete(data_svc.load_data(loop.run_until_complete(data_svc.locate('plugins', dict(enabled=True))))) - loop.run_until_complete(app_svc.load_plugin_expansions(loop.run_until_complete(data_svc.locate('plugins', dict(enabled=True))))) + loop.run_until_complete( + data_svc.load_data( + loop.run_until_complete(data_svc.locate("plugins", dict(enabled=True))) + ) + ) + loop.run_until_complete( + app_svc.load_plugin_expansions( + loop.run_until_complete(data_svc.locate("plugins", dict(enabled=True))) + ) + ) loop.run_until_complete(RestApi(services).enable()) loop.run_until_complete(auth_svc.set_login_handlers(services)) loop.create_task(app_svc.start_sniffer_untrusted_agents()) @@ -71,84 +83,134 @@ def run_tasks(services, run_vue_server=False): loop.create_task(learning_svc.build_model()) loop.create_task(app_svc.watch_ability_files()) loop.run_until_complete(start_server()) - loop.run_until_complete(event_svc.fire_event(exchange='system', queue='ready')) + loop.run_until_complete(event_svc.fire_event(exchange="system", queue="ready")) if run_vue_server: loop.run_until_complete(start_vue_dev_server()) try: - logging.info('All systems ready.') + logging.info("All systems ready.") logging.info(ASCII_BANNER) loop.run_forever() except KeyboardInterrupt: - loop.run_until_complete(services.get('app_svc').teardown(main_config_file=args.environment)) + loop.run_until_complete( + services.get("app_svc").teardown(main_config_file=args.environment) + ) def init_swagger_documentation(app): """Makes swagger documentation available at /api/docs for any endpoints marked for aiohttp_apispec documentation. """ - warnings.filterwarnings( - "ignore", - message="Multiple schemas resolved to the name" - ) + warnings.filterwarnings("ignore", message="Multiple schemas resolved to the name") aiohttp_apispec.setup_aiohttp_apispec( app=app, - title='Caldera', + title="Caldera", version=version.get_version(), - swagger_path='/api/docs', - url='/api/docs/swagger.json', - static_path='/static/swagger' + swagger_path="/api/docs", + url="/api/docs/swagger.json", + static_path="/static/swagger", ) app.middlewares.append(apispec_request_validation_middleware) app.middlewares.append(validation_middleware) async def enable_cors(request, response): - response.headers['Access-Control-Allow-Origin'] = 'http://' + args.uiDevHost + ':3000' - response.headers['Access-Control-Allow-Credentials'] = 'true' - response.headers['Access-Control-Allow-Methods'] = 'GET, POST, PUT, PATCH, DELETE, OPTIONS, HEAD' - response.headers['Access-Control-Allow-Headers'] = 'Access-Control-Allow-Headers, Origin, Accept, X-Requested-With, Content-Type, Access-Control-Request-Method, Access-Control-Request-Headers' + response.headers["Access-Control-Allow-Origin"] = ( + "http://" + args.uiDevHost + ":3000" + ) + response.headers["Access-Control-Allow-Credentials"] = "true" + response.headers["Access-Control-Allow-Methods"] = ( + "GET, POST, PUT, PATCH, DELETE, OPTIONS, HEAD" + ) + response.headers["Access-Control-Allow-Headers"] = ( + "Access-Control-Allow-Headers, Origin, Accept, X-Requested-With, Content-Type, Access-Control-Request-Method, Access-Control-Request-Headers" + ) async def start_vue_dev_server(): await asyncio.create_subprocess_shell( - 'npm run dev', - stdout=sys.stdout, - stderr=sys.stderr, - cwd='./plugins/magma/') - logging.info('VueJS development server is live.') + "npm run dev", stdout=sys.stdout, stderr=sys.stderr, cwd="./plugins/magma/" + ) + logging.info("VueJS development server is live.") -if __name__ == '__main__': +def _get_parser(): + def list_str(values): - return values.split(',') - sys.path.append('') - parser = argparse.ArgumentParser('Welcome to the system') - parser.add_argument('-E', '--environment', required=False, default='local', help='Select an env. file to use') - parser.add_argument("-l", "--log", dest="logLevel", choices=['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'], - help="Set the logging level", default='INFO') - parser.add_argument('--fresh', action='store_true', required=False, default=False, - help='remove object_store on start') - parser.add_argument('-P', '--plugins', required=False, default=os.listdir('plugins'), - help='Start up with a single plugin', type=list_str) - parser.add_argument('--insecure', action='store_true', required=False, default=False, - help='Start caldera with insecure default config values. Equivalent to "-E default".') - parser.add_argument('--uidev', dest='uiDevHost', help='Start VueJS dev server for front-end alongside the caldera server. Provide hostname, i.e. localhost.') - parser.add_argument('--build', action='store_true', required=False, default=False, help='Build the VueJS front-end to serve it from the caldera server.') + return values.split(",") + parser = argparse.ArgumentParser("Welcome to the system") + parser.add_argument( + "-E", + "--environment", + required=False, + default="local", + help="Select an env. file to use", + ) + parser.add_argument( + "-l", + "--log", + dest="logLevel", + choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], + help="Set the logging level", + default="INFO", + ) + parser.add_argument( + "--fresh", + action="store_true", + required=False, + default=False, + help="remove object_store on start", + ) + parser.add_argument( + "-P", + "--plugins", + required=False, + default=os.listdir("plugins"), + help="Start up with a single plugin", + type=list_str, + ) + parser.add_argument( + "--insecure", + action="store_true", + required=False, + default=False, + help='Start caldera with insecure default config values. Equivalent to "-E default".', + ) + parser.add_argument( + "--uidev", + dest="uiDevHost", + help="Start VueJS dev server for front-end alongside the caldera server. Provide hostname, i.e. localhost.", + ) + parser.add_argument( + "--build", + action="store_true", + required=False, + default=False, + help="Build the VueJS front-end to serve it from the caldera server.", + ) + return parser + + +if __name__ == '__main__': + sys.path.append("") + + parser = _get_parser() args = parser.parse_args() setup_logger(getattr(logging, args.logLevel)) if args.insecure: - logging.warning('--insecure flag set. Caldera will use the default.yml config file.') - args.environment = 'default' - elif args.environment == 'local': + logging.warning( + "--insecure flag set. Caldera will use the default.yml config file." + ) + args.environment = "default" + elif args.environment == "local": ensure_local_config() - main_config_path = 'conf/%s.yml' % args.environment - BaseWorld.apply_config('main', BaseWorld.strip_yml(main_config_path)[0]) - logging.info('Using main config from %s' % main_config_path) - BaseWorld.apply_config('agents', BaseWorld.strip_yml('conf/agents.yml')[0]) - BaseWorld.apply_config('payloads', BaseWorld.strip_yml('conf/payloads.yml')[0]) + main_config_path = "conf/%s.yml" % args.environment + BaseWorld.apply_config("main", BaseWorld.strip_yml(main_config_path)[0]) + logging.info("Using main config from %s" % main_config_path) + BaseWorld.apply_config("agents", BaseWorld.strip_yml("conf/agents.yml")[0]) + BaseWorld.apply_config("payloads", BaseWorld.strip_yml("conf/payloads.yml")[0]) data_svc = DataService() knowledge_svc = KnowledgeService() @@ -158,7 +220,7 @@ def list_str(values): Executor, Agent, Link, - AppConfigGlobalVariableIdentifier + AppConfigGlobalVariableIdentifier, ] ) rest_svc = RestService() @@ -167,25 +229,32 @@ def list_str(values): learning_svc = LearningService() event_svc = EventService() - app_svc = AppService(application=web.Application(client_max_size=5120**2, middlewares=[pass_option_middleware])) - app_svc.register_subapp('/api/v2', app.api.v2.make_app(app_svc.get_services())) + app_svc = AppService( + application=web.Application( + client_max_size=5120**2, middlewares=[pass_option_middleware] + ) + ) + app_svc.register_subapp("/api/v2", app.api.v2.make_app(app_svc.get_services())) init_swagger_documentation(app_svc.application) - if (args.uiDevHost): + if args.uiDevHost: if not os.path.exists("./plugins/magma/dist"): logging.info("Building VueJS front-end.") - subprocess.run(['npm', 'run', 'build'], cwd='plugins/magma', check=True) + subprocess.run(["npm", "run", "build"], cwd="plugins/magma", check=True) logging.info("VueJS front-end build complete.") app_svc.application.on_response_prepare.append(enable_cors) - if (args.build): + if args.build: logging.info("Building VueJS front-end.") - subprocess.run(['npm', 'install'], cwd='plugins/magma', check=True) - subprocess.run(['npm', 'run', 'build'], cwd='plugins/magma', check=True) + subprocess.run(["npm", "install"], cwd="plugins/magma", check=True) + subprocess.run(["npm", "run", "build"], cwd="plugins/magma", check=True) logging.info("VueJS front-end build complete.") if args.fresh: - logging.info("Fresh startup: resetting server data. See %s directory for data backups.", DATA_BACKUP_DIR) + logging.info( + "Fresh startup: resetting server data. See %s directory for data backups.", + DATA_BACKUP_DIR, + ) asyncio.get_event_loop().run_until_complete(data_svc.destroy()) asyncio.get_event_loop().run_until_complete(knowledge_svc.destroy()) - run_tasks(services=app_svc.get_services(), run_vue_server=args.uiDevHost) + run_tasks(services=app_svc.get_services(), run_vue_server=args.uiDevHost) \ No newline at end of file