Merge pull request #53 from devplayer55221/firstcontribution

Added exception handling & email param in cloudflare DNS module
PhonePe · Oct 16, 2024 · 00080ba · 00080ba
2 parents d3d21ac + b9f7bab
commit 00080ba
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 12 deletions.
diff --git a/mantis/models/args_model.py b/mantis/models/args_model.py
@@ -21,4 +21,5 @@ class ArgsModel(BaseModel):
     subdomain: str = Field(None)
     list_: bool = False
     list_orgs: bool = False
+    in_scope: bool = False
 
diff --git a/mantis/modules/dns/Cloudflare.py b/mantis/modules/dns/Cloudflare.py
@@ -1,5 +1,6 @@
 import os
 import logging
+from mantis.utils.tool_utils import get_assets_grouped_by_type
 from mantis.config_parsers.config_client import ConfigProvider
 from mantis.tool_base_classes.baseScanner import BaseScanner
 from mantis.models.args_model import ArgsModel
@@ -19,6 +20,7 @@ class Cloudflare(BaseScanner):
 
     async def init(self, args: ArgsModel):
         self.args = args
+        self.db_assets = await get_assets_grouped_by_type(self, args, ASSET_TYPE_TLD)
         return [(self, "Cloudflare")]
 
     async def execute(self, tooltuple):
@@ -32,17 +34,22 @@ async def main(self):
         Prerequisite for this script - A CLoudflare DNS Zone Read only API key \n
         
         """
-        token = None
-
-        cf = CloudFlare.CloudFlare(token, raw=True)
+        token = None    #Edit the value with actual token
         per_page = 100
-
-        zones = cf.zones.get(params={'per_page': per_page, 'page': 0})
         output_dict_list = []
         results = {}
         results["success"] = 0
         results["failure"] = 0
         try:
+            try:
+                cf = CloudFlare.CloudFlare("", token, raw=True)
+                zones = cf.zones.get(params={'per_page': per_page, 'page': 0})
+                results["success"] += 1
+            except Exception as e:
+                results["failure"] += 1
+                results["exception"] = str(e)
+                logging.error("[!] Error - {}".format(str(e)))
+
             for zone_page in range(zones['result_info']['total_pages']):
                 zones = cf.zones.get(params={'per_page': per_page, 'page': zone_page})
                 for zone in zones['result']:
@@ -53,14 +60,28 @@ async def main(self):
                     records = cf.zones.dns_records.get(zone['id'], params={'per_page': per_page, 'page': record_page})['result']
                     for record in records:
                         domain_dict = {}
-                        domain_dict['_id'] = record['name']
-                        domain_dict['asset'] = record['name']
-                        if AssetType.check_tld(record['name']):
-                            domain_dict['asset_type'] = ASSET_TYPE_TLD
+                        if(self.args.in_scope == True):
+                            print(self.db_assets)
+                            for asset in self.db_assets:
+                                if(asset in record['name']):
+                                    domain_dict['_id'] = record['name']
+                                    domain_dict['asset'] = record['name']
+                                    if AssetType.check_tld(record['name']):
+                                        domain_dict['asset_type'] = ASSET_TYPE_TLD
+                                    else:
+                                        domain_dict['asset_type'] = ASSET_TYPE_SUBDOMAIN
+                                    domain_dict['org'] = self.args.org
+                                    output_dict_list.append(domain_dict)
+                                    break
                         else:
-                            domain_dict['asset_type'] = ASSET_TYPE_SUBDOMAIN
-                        domain_dict['org'] = self.args.org
-                        output_dict_list.append(domain_dict)
+                            domain_dict['_id'] = record['name']
+                            domain_dict['asset'] = record['name']
+                            if AssetType.check_tld(record['name']):
+                                domain_dict['asset_type'] = ASSET_TYPE_TLD
+                            else:
+                                domain_dict['asset_type'] = ASSET_TYPE_SUBDOMAIN
+                            domain_dict['org'] = self.args.org
+                            output_dict_list.append(domain_dict)
             await CrudUtils.insert_assets(output_dict_list, source='internal')
             results["success"] = 1
             return results

diff --git a/mantis/utils/args_parse.py b/mantis/utils/args_parse.py
@@ -229,6 +229,12 @@ def args_parse() -> ArgsModel:
         scan_parser.add_argument('--sub', 
                                  dest = 'subdomain',
                                  help='Subdomain to scan')
+
+        scan_parser.add_argument('-is', '--in_scope', 
+                            dest = 'in_scope',
+                            help = 'List only the records from nameserver that are in scope',
+                            action = 'store_true' 
+                            )
 
 
         list_parser = subparser.add_parser("list", help="List entities present in db", usage=ArgsParse.list_msg())