From 0c0ef9d8d515745ac6ad0a0d28ff028cd1de1d3a Mon Sep 17 00:00:00 2001 From: Marcus Bakker Date: Sat, 13 Feb 2021 13:27:51 +0100 Subject: [PATCH] Improved exception handling for MITRE's CTI TAXII server connection problems --- editor/script/refresh_attack_data.py | 3 ++- editor/script/requirements.txt | 3 ++- generic.py | 3 ++- 3 files changed, 6 insertions(+), 3 deletions(-) diff --git a/editor/script/refresh_attack_data.py b/editor/script/refresh_attack_data.py index 26685c2..3a6b38c 100644 --- a/editor/script/refresh_attack_data.py +++ b/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() diff --git a/editor/script/requirements.txt b/editor/script/requirements.txt index 39510eb..60ceb84 100644 --- a/editor/script/requirements.txt +++ b/editor/script/requirements.txt @@ -1 +1,2 @@ -attackcti==0.3.3 \ No newline at end of file +attackcti==0.3.3 +taxii2-client==2.2.2 \ No newline at end of file diff --git a/generic.py b/generic.py index a649024..375560a 100644 --- a/generic.py +++ b/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()