Merge pull request #35 from PhonePe/feature/list_orgs

Added list subcommand

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') 

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
 

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")

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

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)

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))