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

master
Marcus Bakker 2021-02-13 13:27:51 +01:00
parent bd8777f7b6
commit 0c0ef9d8d5
3 changed files with 6 additions and 3 deletions

View File

@ -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()

View File

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

View File

@ -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()