Added a graceful quit when the MITRE's CTI server cannot be reached

editor/script/refresh_attack_data.py

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

generic.py

@@ -8,6 +8,7 @@ from ruamel.yaml.timestamp import TimeStamp as ruamelTimeStamp
 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
 
 # Due to performance reasons the import of attackcti is within the function that makes use of this library.
 
@@ -54,8 +55,11 @@ def load_attack_data(data_type):
                 if not (dt.now() - write_time).total_seconds() >= EXPIRE_TIME:
                     # the first item in the list contains the ATT&CK data
                     return cached[0]
-
+        try:
-        mitre = attack_client()
+            mitre = attack_client()
+        except exceptions.ConnectionError:
+            print("[!] Cannot connect to MITRE's CTI TAXII server")
+            quit()
 
     attack_data = None
     if data_type == DATA_TYPE_STIX_ALL_RELATIONSHIPS: