diff --git a/mantis/config_parsers/logging_utils.py b/mantis/config_parsers/logging_utils.py index 03df579..5e710af 100644 --- a/mantis/config_parsers/logging_utils.py +++ b/mantis/config_parsers/logging_utils.py @@ -36,10 +36,10 @@ def configure_logging(args: ArgsModel): log_config = ConfigProvider.get_config().logging_debug logging.config.dictConfig(log_config) logging.StreamHandler().setFormatter(CustomFormatter()) - logging.info('MANTIS ASSET DISCOVERY - STARTED') + logging.info('MANTIS Workflow - STARTED') logging.info("Debug mode enabled") else: log_config = ConfigProvider.get_config().logging logging.config.dictConfig(log_config) logging.StreamHandler().setFormatter(CustomFormatter()) - logging.info('MANTIS ASSET DISCOVERY - STARTED') + logging.info('MANTIS Workflow - STARTED') diff --git a/mantis/models/args_model.py b/mantis/models/args_model.py index 82f6c83..a843f94 100644 --- a/mantis/models/args_model.py +++ b/mantis/models/args_model.py @@ -19,4 +19,6 @@ class ArgsModel(BaseModel): verbose: bool = False thread_count: int = 3 subdomain: str = Field(None) + list_: bool = False + list_orgs: bool = False \ No newline at end of file diff --git a/mantis/utils/args_parse.py b/mantis/utils/args_parse.py index 9b2d19d..43916c2 100644 --- a/mantis/utils/args_parse.py +++ b/mantis/utils/args_parse.py @@ -44,7 +44,14 @@ def scan_msg(name=None): \033[0;32mmantis scan -o example_org -a example_app\033[0m ''' + @staticmethod + def list_msg(name=None): + return ''' + \033[1;34mList:\033[0m + \033[0;32mmantis list {subcommand}\033[0m + ''' + @staticmethod def args_parse() -> ArgsModel: parsed_args = {} @@ -56,7 +63,7 @@ def args_parse() -> ArgsModel: usage=ArgsParse.msg() ) - subparser = parser.add_subparsers(title="subparser", dest="subcommand") + subparser = parser.add_subparsers(title="TYPE", dest="subcommand") onboard_parser = subparser.add_parser("onboard", help="Onboard a target", usage=ArgsParse.onboard_msg()) @@ -224,6 +231,12 @@ def args_parse() -> ArgsModel: help='Subdomain to scan') + list_parser = subparser.add_parser("list", help="List entities present in db", usage=ArgsParse.list_msg()) + + list_sub_parser = list_parser.add_subparsers(title="List Subcommands", dest="list_sub_command") + + list_sub_parser.add_parser("orgs", help="List orgs present in DB") + # display help, if no arguments are passed args = parser.parse_args(args=None if argv[1:] else ['--help']) logging.info(f"Arguments Passed - {args}") @@ -240,44 +253,46 @@ def args_parse() -> ArgsModel: parsed_args['input_type'] = "file" parsed_args['input'] = str(args.file_name) - if args.aws_profiles: - parsed_args["aws_profiles"] = args.aws_profiles.split(',') - else: - parsed_args["aws_profiles"] = ['default'] + if args.subcommand != "list": - if args.workflow: - parsed_args['workflow'] = args.workflow - else: - parsed_args['workflow'] = 'default' + if args.aws_profiles: + parsed_args["aws_profiles"] = args.aws_profiles.split(',') + else: + parsed_args["aws_profiles"] = ['default'] - parsed_args['org'] = args.org + if args.workflow: + parsed_args['workflow'] = args.workflow + else: + parsed_args['workflow'] = 'default' - if args.app: - parsed_args["app"] = args.app - - if args.passive: - parsed_args["passive"] = True - - if args.stale: - parsed_args["stale"] = True + parsed_args['org'] = args.org - if args.ignore_stale: - parsed_args["ignore_stale"] = True - - if args.use_ray: - parsed_args["use_ray"] = True - - if args.num_actors: - parsed_args["num_actors"] = args.num_actors - - if args.delete_logs: - parsed_args["delete_logs"] = args.delete_logs + if args.app: + parsed_args["app"] = args.app + + if args.passive: + parsed_args["passive"] = True + + if args.stale: + parsed_args["stale"] = True + + if args.ignore_stale: + parsed_args["ignore_stale"] = True + + if args.use_ray: + parsed_args["use_ray"] = True + + if args.num_actors: + parsed_args["num_actors"] = args.num_actors + + if args.delete_logs: + parsed_args["delete_logs"] = args.delete_logs - if args.verbose: - parsed_args["verbose"] = True + if args.verbose: + parsed_args["verbose"] = True - if args.thread_count: - parsed_args["thread_count"] = args.thread_count + if args.thread_count: + parsed_args["thread_count"] = args.thread_count if args.subcommand == "scan": if args.subdomain: @@ -287,6 +302,12 @@ def args_parse() -> ArgsModel: if args.subdomain: parsed_args["subdomain"] = args.host + if args.subcommand == "list": + parsed_args["list_"] = True + + if args.list_sub_command == "orgs": + parsed_args["list_orgs"] = True + args_pydantic_obj = ArgsModel.parse_obj(parsed_args) logging.info(f'parsed args - {args_pydantic_obj}') logging.info(f"Parsing Arguements - Completed") diff --git a/mantis/utils/list_subcommand_utils.py b/mantis/utils/list_subcommand_utils.py new file mode 100644 index 0000000..243e953 --- /dev/null +++ b/mantis/utils/list_subcommand_utils.py @@ -0,0 +1,21 @@ +from mantis.db.crud_assets import read_assets + +async def get_orgs(): + pipeline = [] + + pipeline.extend([ + {"$group": { + "_id" : None, + "orgs":{ + "$addToSet": "$org" + } + } + } + ]) + orgs = await read_assets(pipeline) + + if orgs: + return orgs[0]["orgs"] + + else: + return None \ No newline at end of file diff --git a/mantis/workflows/list_workflow.py b/mantis/workflows/list_workflow.py new file mode 100644 index 0000000..df0c157 --- /dev/null +++ b/mantis/workflows/list_workflow.py @@ -0,0 +1,17 @@ +import logging +from mantis.models.args_model import ArgsModel +from mantis.utils.list_subcommand_utils import get_orgs + +class ListWorkflow: + + @staticmethod + async def executor(args: ArgsModel): + + if args.list_orgs: + logging.info("Getting orgs from database") + orgs = await get_orgs() + print() + print(f'''Total Number of orgs onboarded: {len(orgs)}''') + print() + for org in orgs: + print(org) \ No newline at end of file diff --git a/mantis/workflows/mantis_workflow.py b/mantis/workflows/mantis_workflow.py index cd3a40a..5befd7c 100644 --- a/mantis/workflows/mantis_workflow.py +++ b/mantis/workflows/mantis_workflow.py @@ -1,10 +1,14 @@ from mantis.models.args_model import ArgsModel from mantis.modules.workflow import Workflow +from mantis.workflows.list_workflow import ListWorkflow import asyncio class MantisWorkflow: @staticmethod def select_workflow(args: ArgsModel) -> None: - asyncio.run(Workflow.workflow_executor(args)) + if args.list_: + asyncio.run(ListWorkflow.executor(args)) + else: + asyncio.run(Workflow.workflow_executor(args)) \ No newline at end of file