Improved exception handling for MITRE's CTI TAXII server connection problems

editor/script/refresh_attack_data.py

@@ -1,6 +1,7 @@
 import json
 from attackcti import attack_client
 from requests import exceptions
+from stix2 import datastore
 
 FILE_DATA_SOURCES = 'data_sources.json'
 FILE_TECHNIQUES = 'techniques.json'
@@ -16,7 +17,7 @@ class ATTACKData():
     def __init__(self):
         try:
             self.mitre = attack_client()
-        except exceptions.ConnectionError:
+        except (exceptions.ConnectionError, datastore.DataSourceError):
             print("[!] Cannot connect to MITRE's CTI TAXII server")
             quit()
         self.attack_cti_techniques = self.mitre.get_enterprise_techniques()

editor/script/requirements.txt

@@ -1 +1,2 @@
 attackcti==0.3.3
+taxii2-client==2.2.2

generic.py

@@ -9,6 +9,7 @@ from upgrade import upgrade_yaml_file, check_yaml_updated_to_sub_techniques
 from constants import *
 from health import check_yaml_file_health
 from requests import exceptions
+from stix2 import datastore
 
 # Due to performance reasons the import of attackcti is within the function that makes use of this library.
 
@@ -57,7 +58,7 @@ def load_attack_data(data_type):
                     return cached[0]
         try:
             mitre = attack_client()
-        except exceptions.ConnectionError:
+        except (exceptions.ConnectionError, datastore.DataSourceError):
             print("[!] Cannot connect to MITRE's CTI TAXII server")
             quit()