From 701516e4810a95400767511bf19f3f6ca5ca1bc4 Mon Sep 17 00:00:00 2001
From: Marcus Bakker <marcus@mbsecure.nl>
Date: Wed, 10 Feb 2021 13:17:46 +0100
Subject: [PATCH] Added a graceful quit when the MITRE's CTI server cannot be
 reached

---
 editor/script/refresh_attack_data.py | 7 ++++++-
 generic.py                           | 8 ++++++--
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/editor/script/refresh_attack_data.py b/editor/script/refresh_attack_data.py
index ec00bfb..26685c2 100644
--- a/editor/script/refresh_attack_data.py
+++ b/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)
diff --git a/generic.py b/generic.py
index 81c10b8..a649024 100644
--- a/generic.py
+++ b/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]
-
-        mitre = attack_client()
+        try:
+            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: