SixGill and Shodan Custom Exceptions

build/lib/pe_reports/exceptionTest.py

@@ -0,0 +1,3 @@
+from pe_reports.peExcept import SixGillAlertInsertException as ok
+
+raise ok("test1","test2")

src/pe_reports/helpers/__init__.py

@@ -0,0 +1,2 @@
+"""Init file for module implementation."""
+

src/pe_reports/helpers/bulletin/__init__.py

@@ -0,0 +1,2 @@
+"""Init file for module implementation."""
+

src/pe_reports/peExcept.py

@@ -0,0 +1,50 @@
+from pe_reports import CENTRAL_LOGGING_FILE
+import logging
+LOGGER = logging.getLogger(__name__)
+
+class CustomException(Exception):
+    """No Data in the domain list. """
+    def __init__(self, org, message="DnsTwist data did not collect seccessfully for org"):
+        self.org = org
+        self.message = message
+        self.log()
+        super().__init__(self.message)
+    def __str__(self):
+        return f'{self.message} : {self.org}'
+    def log(self):
+        LOGGER.error(f'{self.message} : {self.org}')
+
+
+class SixGillApiException(Exception):
+    """Exception for failure in sixgill alert insertions. """
+    def __init__(self, org_id,six_gill_id, message="CyberSixGill was call was unsuccessful"):
+        self.six_gill_id = six_gill_id  
+        self.org = org
+        self.message = message
+        super().__init__(self.message)
+    def __str__(self):
+        return f'{self.message}, Org ID : {self.org}, Six Gill ID: {self.six_gill_id}'
+    def log(self):
+        LOGGER.error(f'{self.message} : {self.org}, Six Gill ID: {self.six_gill_id}')
+
+class SixGillDatabaseException(Exception):
+    """Exception for failure in sixgill alert insertions. """
+    def __init__(self, org_id, message="Data insertion into the database was unsuccessful"):
+        self.org = org
+        self.message = message
+        super().__init__(self.message)
+    def __str__(self):
+        return f'{self.message} , Source UID : {self.org}'
+    def log(self):
+        LOGGER.error(f'{self.message} : {self.org}')
+
+class ShodanIPFailure(Exception):
+    """Exception for failing to lookup Shodan ips . """
+    def __init__(self, thread_name, org_name):
+        self.org = org
+        self.message = message
+        super().__init__(self.message)
+    def __str__(self):
+        return "{} Failed fetching IPs for {}.".format(thread_name, org_name)
+    def log(self):
+        LOGGER.error(f'{self.message} : {self.org}')

src/pe_source/cybersixgill.py

@@ -6,6 +6,7 @@
 
 # cisagov Libraries
 from pe_reports import app
+from pe_reports.peExcept import SixGillApiException, SixGillDatabaseException
 
 from .data.pe_db.db_query_source import (
     get_breaches,
@@ -128,16 +129,19 @@ def get_alerts(self, org_id, sixgill_org_id, pe_org_uid, source_uid):
             alerts_df["data_source_uid"] = source_uid
             # Rename columns
             alerts_df = alerts_df.rename(columns={"id": "sixgill_id"})
-        except Exception as e:
-            LOGGER.error("Failed fetching alert data for %s", org_id)
+        except SixGillApiException(
+            org_id, sixgill_org_id, "Failed Fetching Alert Data"
+        ) as e:
+            # LOGGER.error("Failed fetching alert data for %s", org_id)
             LOGGER.error(e)
             return 1
 
         # Insert alert data into the PE database
         try:
             insert_sixgill_alerts(alerts_df)
-        except Exception as e:
-            LOGGER.error("Failed inserting alert data for %s", org_id)
+        except SixGillApiException(
+            org_id, sixgill_org_id, "Failed inserting alert data"
+        ) as e:
             LOGGER.error(e)
             return 1
         return 0
@@ -149,8 +153,10 @@ def get_mentions(self, org_id, sixgill_org_id, pe_org_uid, source_uid):
         # Fetch org aliases from Cybersixgill
         try:
             aliases = alias_organization(sixgill_org_id)
-        except Exception as e:
-            LOGGER.error("Failed fetching aliases for %s", org_id)
+        except SixGillApiException(
+            org_id, sixgill_org_id, "Failed fetching aliases"
+        ) as e:
+            # LOGGER.error("Failed fetching aliases for %s", org_id)
             LOGGER.error(e)
             return 1
 
@@ -161,15 +167,19 @@ def get_mentions(self, org_id, sixgill_org_id, pe_org_uid, source_uid):
             mentions_df["organizations_uid"] = pe_org_uid
             # Add data source uid
             mentions_df["data_source_uid"] = source_uid
-        except Exception as e:
-            LOGGER.error("Failed fetching mentions for %s", org_id)
+        except SixGillApiException(
+            org_id, sixgill_org_id, "Failed fetching mentions"
+        ) as e:
+            # LOGGER.error("Failed fetching mentions for %s", org_id)
             LOGGER.error(e)
             return 1
 
         # Insert mention data into the PE database
         try:
             insert_sixgill_mentions(mentions_df)
-        except Exception as e:
+        except SixGillApiException(
+            org_id, sixgill_org_id, "Failed Inserting Alert Data"
+        ) as e:
             LOGGER.error("Failed inserting mentions for %s", org_id)
             LOGGER.error(e)
             return 1
@@ -182,8 +192,10 @@ def get_credentials(self, org_id, sixgill_org_id, pe_org_uid, source_uid):
         # Fetch org root domains from Cybersixgill
         try:
             roots = root_domains(sixgill_org_id)
-        except Exception as e:
-            LOGGER.error("Failed fetching root domains for %s", org_id)
+        except SixGillApiException(
+            org_id, sixgill_org_id, "Failed fetching root domains"
+        ) as e:
+            # LOGGER.error("Failed fetching root domains for %s", org_id)
             LOGGER.error(e)
             return 1
 
@@ -193,8 +205,10 @@ def get_credentials(self, org_id, sixgill_org_id, pe_org_uid, source_uid):
             creds_df["organizations_uid"] = pe_org_uid
             # Add data source uid
             creds_df["data_source_uid"] = source_uid
-        except Exception as e:
-            LOGGER.error("Failed fetching credentials for %s", org_id)
+        except SixGillApiException(
+            org_id, sixgill_org_id, "Failed fetching credentials"
+        ) as e:
+            # LOGGER.error("Failed fetching credentials for %s", org_id)
             LOGGER.error(e)
             return 1
 
@@ -234,8 +248,10 @@ def get_credentials(self, org_id, sixgill_org_id, pe_org_uid, source_uid):
         # Insert breach data into the PE database
         try:
             insert_sixgill_breaches(creds_breach_df)
-        except Exception as e:
-            LOGGER.error("Failed inserting breaches for %s", org_id)
+        except SixGillDatabaseException(
+            org_id, sixgill_org_id, "Failed inserting data"
+        ) as e:
+            # LOGGER.error("Failed inserting breaches for %s", org_id)
             LOGGER.error(e)
             return 1
 
@@ -267,8 +283,10 @@ def get_credentials(self, org_id, sixgill_org_id, pe_org_uid, source_uid):
         ]
         try:
             insert_sixgill_credentials(creds_df)
-        except Exception as e:
-            LOGGER.error("Failed inserting credentials for %s", org_id)
+        except SixGillDatabaseException(
+            org_id, sixgill_org_id, "Failed Inserting Alert Data"
+        ) as e:
+            # LOGGER.error("Failed inserting credentials for %s", org_id)
             LOGGER.error(e)
             return 1
         return 0
@@ -293,16 +311,18 @@ def get_topCVEs(self, source_uid):
                 except Exception:
                     summary = ""
                 top_cve_df.at[index, "summary"] = summary
-        except Exception as e:
-            LOGGER.error("Failed fetching top CVEs.")
+        except SixGillApiException(source_uid, 0, "Failed fetching CVEs") as e:
+            # LOGGER.error("Failed fetching top CVEs.")
             LOGGER.error(e)
             return 1
 
         # Insert credential data into the PE database
         try:
             insert_sixgill_topCVEs(top_cve_df)
-        except Exception as e:
-            LOGGER.error("Failed inserting top CVEs.")
+        except SixGillDatabaseException(
+            source_uid, 0, "Failed inserting top CVEs"
+        ) as e:
+            # LOGGER.error("Failed inserting top CVEs.")
             LOGGER.error(e)
             return 1
         return 0

src/pe_source/data/shodan/shodan_search.py

@@ -11,6 +11,7 @@
 
 # cisagov Libraries
 from pe_reports import app
+from pe_reports.peExcept import ShodanIPFailure
 from pe_source.data.pe_db.db_query_source import (
     get_data_source_uid,
     get_ips,
@@ -31,9 +32,8 @@ def run_shodan_thread(api, org_chunk, thread_name):
         start, end = get_dates()
         try:
             ips = get_ips(org_uid)
-        except Exception as e:
-            LOGGER.error("{} Failed fetching IPs for {}.".format(thread_name, org_name))
-            LOGGER.error("{} {} - {}".format(thread_name, e, org_name))
+        except ShodanIPFailure(thread_name, org_name) as e:
+            LOGGER.error(e)
             failed.append("{} fetching IPs".format(org_name))
             continue