---
title: 'IntelliSIEM and Test-Driven Development: Creating Robust Error Handling and Log Rotation'
url: 'https://bsb.perryconsulting.net/blog/intellisiem-and-test-driven-development-creating-robust-error-handling-and-log-rotation'
markdown: 'https://bsb.perryconsulting.net/blog/intellisiem-and-test-driven-development-creating-robust-error-handling-and-log-rotation.md'
date: '2024-10-01'
description: "Previously, we implemented a basic testing framework to start handling data errors, corrupt data, and edge cases. In Part 4, we're going to kick it up a notch. It's time to implement robust error handling for each discrete task! We left off having implemented a basic testing framework to drive code implementation for edge cases and errors. Now we're going to build up that framework by developing tests specifically for each discrete function and attempting to ensure we have solid error handling for anything that comes our way. We're also going to implement log rotation while we're at it, to ensure that there's a reasonable amount of logs captured without taking up too much drive space. Let's get going! To see this project in its curr"
taxonomy:
category:
- projects
tag:
- python
- security
archives_month:
- oct_2024
archives_year:
- '2024'
---
# IntelliSIEM and Test-Driven Development: Creating Robust Error Handling and Log Rotation
01 OCT 2024 [python](https://bsb.perryconsulting.net/blog/tag:python) [security](https://bsb.perryconsulting.net/blog/tag:security)
*Previously, we implemented a basic testing framework to start handling data errors, corrupt data, and edge cases. In Part 4, we're going to* *kick it up a notch. It's time to implement robust error handling for each discrete task*!
We left off having implemented a basic testing framework to drive code implementation for edge cases and errors. Now we're going to build up that framework by developing tests specifically for each discrete function and attempting to ensure we have solid error handling for anything that comes our way. We're also going to implement log rotation while we're at it, to ensure that there's a reasonable amount of logs captured without taking up too much drive space. Let's get going!
To see this project in its current state, head over to my repo and check it out: [IntelliSIEM on GitHub](https://github.com/perryconsulting/IntelliSIEM)
---
### **Writing Better Tests To Drive Implementation**
One of the unwritten goals of this project is to develop the application while adhering to Test-Driven Development (TDD) principles as much as possible. This means writing tests first and then modifying our code to meet the expectations of those tests. One area where TDD is particularly effective is in implementing robust error handling. In this stage, the aim is to capture as many possible error scenarios in the data collection process as possible, such as invalid data types, missing keys, corrupt data, or empty responses. With those goals in mind, we need to refactor our tests, breaking related tests into their own files and adding more discrete tests to define better how the code should work. This can be a somewhat tedious process, so in the interest of not putting you through it all, here is the newly refactored test suite:
PYTHON
```
1# IntelliSIEM Copyright 2024, Rob Perry
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15import pytest
16import requests
17from src.api_client import APIClient, APIError
18
19def test_api_client_success(mocker):
20 """
21 Test API client handling a successful request.
22 This ensures that valid data is returned when the API call is successful.
23 """
24 mock_response = [{'id': 'test_data'}]
25 mocker.patch('requests.get',
26 return_value=mocker.Mock(status_code=200, json=mocker.Mock(return_value=mock_response)))
27 client = APIClient("https://example.com/api", {"Authorization": "Bearer token"})
28 result = client.get_data("endpoint")
29 assert result == mock_response # Expecting mock response data
30
31def test_api_client_timeout(mocker):
32 """
33 Test API client handling of a timeout error.
34 This ensures that the appropriate exception and message are raised and logged.
35 """
36 mocker.patch('requests.get', side_effect=requests.exceptions.Timeout)
37 client = APIClient("https://example.com/api", {"Authorization": "Bearer token"})
38 with pytest.raises(APIError, match="API request timed out."):
39 client.get_data("endpoint")
40
41def test_api_client_rate_limit(mocker):
42 """
43 Test API client handling of rate limiting.
44 This ensures that the APIError exception is raised when the rate limit is exceeded.
45 """
46 mock_response = mocker.Mock(status_code=429, headers={"Retry-After": "5"})
47 mocker.patch('requests.get', return_value=mock_response)
48 client = APIClient("https://example.com/api", {"Authorization": "Bearer token"})
49 with pytest.raises(APIError, match="Rate limit exceeded."):
50 client.get_data("endpoint")
51
52def test_api_client_invalid_json(mocker):
53 """
54 Test API client handling of invalid JSON response.
55 This test simulates a ValueError when the JSON response is not valid.
56 """
57 mock_response = mocker.Mock(status_code=200, json=mocker.Mock(side_effect=ValueError))
58 mocker.patch('requests.get', return_value=mock_response)
59 client = APIClient("https://example.com/api", {"Authorization": "Bearer token"})
60 with pytest.raises(APIError, match="Invalid response format."):
61 client.get_data("endpoint")
62
63def test_api_client_connection_error(mocker):
64 """
65 Test API client handling of a connection error.
66 This ensures that connection issues are properly caught and logged.
67 """
68 mocker.patch('requests.get', side_effect=requests.ConnectionError)
69 client = APIClient("https://example.com/api", {"Authorization": "Bearer token"})
70 with pytest.raises(APIError, match="Connection error occurred."):
71 client.get_data("endpoint")
72
73def test_api_client_authentication_error(mocker):
74 """
75 Test API client handling of an authentication error.
76 This simulates an authentication failure with a 401 status code.
77 """
78 mock_response = mocker.Mock(status_code=401)
79 mocker.patch('requests.get', return_value=mock_response)
80 client = APIClient("https://example.com/api", {"Authorization": "Invalid token"})
81 with pytest.raises(APIError, match="Authentication failed."):
82 client.get_data("endpoint")
83
84def test_api_client_permission_error(mocker):
85 """
86 Test API client handling of a permission error.
87 This simulates a permission failure with a 403 status code.
88 """
89 mock_response = mocker.Mock(status_code=403)
90 mocker.patch('requests.get', return_value=mock_response)
91 client = APIClient("https://example.com/api", {"Authorization": "Bearer token"})
92 with pytest.raises(APIError, match="Permission denied."):
93 client.get_data("endpoint")
```
PYTHON
```
1# IntelliSIEM Copyright 2024, Rob Perry
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15import pytest
16from src.config import load_config, ConfigError
17
18def test_load_config_success(mocker):
19 """
20 Test successful loading of the configuration file.
21 This ensures that a valid configuration is loaded and returned correctly.
22 """
23 mock_config = """
24 api_keys:
25 alien_vault: "VALID_ALIEN_VAULT_API_KEY"
26 virus_total: "VALID_VIRUS_TOTAL_API_KEY"
27 """
28 mocker.patch('builtins.open', mocker.mock_open(read_data=mock_config))
29 config = load_config()
30 assert config['api_keys']['alien_vault'] == "VALID_ALIEN_VAULT_API_KEY"
31 assert config['api_keys']['virus_total'] == "VALID_VIRUS_TOTAL_API_KEY"
32
33def test_load_config_missing_file():
34 """
35 Test configuration loading with a missing file.
36 This ensures that a FileNotFoundError is raised when the config file does not exist.
37 """
38 with pytest.raises(FileNotFoundError):
39 load_config('config/non_existent_config.yaml')
40
41def test_load_config_invalid_format(mocker):
42 """
43 Test configuration loading with an invalid format.
44 This simulates an invalid YAML format and expects a ConfigError.
45 """
46 invalid_yaml = """
47 api_keys:
48 alien_vault: !!invalid
49 """
50 mocker.patch('builtins.open', mocker.mock_open(read_data=invalid_yaml))
51 with pytest.raises(ConfigError, match="Configuration file error."):
52 load_config()
53
54def test_load_config_permission_error(mocker):
55 """
56 Test configuration loading with a permission error.
57 This simulates a PermissionError when accessing the config file.
58 """
59 mock_open = mocker.patch('builtins.open', mocker.mock_open())
60 mock_open.side_effect = PermissionError
61 with pytest.raises(ConfigError, match="Permission denied for config file."):
62 load_config()
63
64def test_load_config_disk_space_error(mocker):
65 """
66 Test configuration loading with a disk space error.
67 This simulates an OSError due to lack of disk space.
68 """
69 mock_open = mocker.patch('builtins.open', mocker.mock_open())
70 mock_open.side_effect = OSError("No space left on device")
71 with pytest.raises(ConfigError, match="Disk space issue while loading configuration."):
72 load_config()
```
PYTHON
```
1# IntelliSIEM Copyright 2024, Rob Perry
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# tests/test_data_collection.py
15
16import pytest
17from src.data_collection import fetch_threat_data, DataError
18
19def test_fetch_threat_data_success(mocker):
20 """
21 Test successful fetching of threat data.
22 This ensures that data is correctly collected from the APIs and returned.
23 """
24 mock_data_alien = [{'id': 'test1'}]
25 mock_data_virus = [{'id': 'test2'}]
26
27 # Create separate mock instances for each APIClient
28 mock_alien_vault_client = mocker.Mock()
29 mock_alien_vault_client.get_data.return_value = mock_data_alien
30
31 mock_virus_total_client = mocker.Mock()
32 mock_virus_total_client.get_data.return_value = mock_data_virus
33
34 # Patch the APIClient constructor to return the respective mock instances
35 mocker.patch('src.data_collection.APIClient', side_effect=[mock_alien_vault_client, mock_virus_total_client])
36
37 threats = fetch_threat_data()
38
39 # Verify that the mock was called only once per client
40 mock_alien_vault_client.get_data.assert_called_once()
41 mock_virus_total_client.get_data.assert_called_once()
42
43 assert len(threats) == 2 # Expecting two entries: one from each mock data list
44 assert threats == [{'id': 'test1'}, {'id': 'test2'}]
45
46def test_fetch_threat_data_missing_keys(mocker):
47 """
48 Test data collection handling of missing required data keys.
49 This ensures that missing keys in the data trigger the appropriate exception.
50 """
51 # Mock get_data to return an entry missing the required 'id' key
52 mocker.patch('src.api_client.APIClient.get_data', return_value=[{'missing_key': 'value'}])
53 with pytest.raises(DataError, match="Missing required data keys."):
54 fetch_threat_data()
55
56def test_fetch_threat_data_corrupt_data(mocker):
57 """
58 Test data collection handling of corrupt data.
59 This ensures that corrupt data is caught and triggers the DataError exception.
60 """
61 # Mock get_data to return structurally corrupt data with an unexpected type for 'id'
62 mocker.patch('src.api_client.APIClient.get_data',
63 return_value=[{'id': {'unexpected': 'dict'}, 'nested_data': ['unexpected_list']}])
64
65 # Expect DataError to be raised for corrupt data
66 with pytest.raises(DataError, match="Corrupt data received."):
67 fetch_threat_data()
68
69def test_fetch_threat_data_empty_response(mocker):
70 """
71 Test fetching threat data with an empty response.
72 This ensures that the function handles empty responses gracefully.
73 """
74 mocker.patch('src.api_client.APIClient.get_data', return_value=[])
75 threats = fetch_threat_data()
76 assert threats == [] # Expecting an empty list when no data is returned
77
78def test_fetch_threat_data_invalid_data_type(mocker):
79 """
80 Test data collection handling of invalid data type.
81 This simulates an unexpected data type from the API response.
82 """
83 # Mock get_data to return a non-list data type, such as a string
84 mocker.patch('src.api_client.APIClient.get_data', return_value="invalid_data")
85 with pytest.raises(DataError, match="Unexpected data type received."):
86 fetch_threat_data()
```
PYTHON
```
1# IntelliSIEM Copyright 2024, Rob Perry
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15import pytest
16from src.error_handling import APIError, ConfigError, DataError, log_error
17import logging
18
19def test_api_error():
20 """
21 Test that the APIError exception can be raised and caught correctly.
22 """
23 with pytest.raises(APIError, match="API issue occurred"):
24 raise APIError("API issue occurred.")
25
26def test_config_error():
27 """
28 Test that the ConfigError exception can be raised and caught correctly.
29 """
30 with pytest.raises(ConfigError, match="Configuration issue occurred"):
31 raise ConfigError("Configuration issue occurred.")
32
33def test_data_error():
34 """
35 Test that the DataError exception can be raised and caught correctly.
36 """
37 with pytest.raises(DataError, match="Data issue occurred"):
38 raise DataError("Data issue occurred.")
39
40def test_log_error(mocker):
41 """
42 Test that the log_error function logs an error message correctly.
43 """
44 mocker.patch('logging.error')
45 log_error("Test error message.")
46 logging.error.assert_called_once_with("Test error message.")
47
48def test_log_warning(mocker):
49 """
50 Test that the log_error function logs a warning message correctly.
51 """
52 mocker.patch('logging.warning')
53 log_error("Test warning message.", level="WARNING")
54 logging.warning.assert_called_once_with("Test warning message.")
55
56def test_log_info(mocker):
57 """
58 Test that the log_error function logs an informational message correctly.
59 """
60 mocker.patch('logging.info')
61 log_error("Test info message.", level="INFO")
62 logging.info.assert_called_once_with("Test info message.")
```
PYTHON
```
1# IntelliSIEM Copyright 2024, Rob Perry
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15import logging
16import os
17from logging.handlers import RotatingFileHandler
18import pytest
19
20@pytest.fixture
21def setup_rotating_log():
22 """
23 Set up a rotating log handler for testing.
24 """
25 # Define the log file path and cleanup if exists
26 log_file = 'data/test_error.log'
27 for i in range(6):
28 try:
29 os.remove(f"{log_file}.{i}")
30 except FileNotFoundError:
31 pass
32 try:
33 os.remove(log_file)
34 except FileNotFoundError:
35 pass
36
37 # Set up rotating file handler
38 handler = RotatingFileHandler(log_file, maxBytes=200, backupCount=5)
39 handler.setLevel(logging.ERROR)
40 handler.setFormatter(logging.Formatter('%(asctime)s %(levelname)s: %(message)s'))
41 logger = logging.getLogger('test_logger')
42 logger.setLevel(logging.ERROR)
43 logger.addHandler(handler)
44 return logger
45
46def test_log_rotation(setup_rotating_log):
47 """
48 Test log rotation by generating log entries to exceed maxBytes limit.
49 """
50 logger = setup_rotating_log
51 log_message = "This is a test log message to fill up the log file and trigger rotation."
52 log_file = 'data/test_error.log'
53
54 # Write multiple log entries to exceed the maxBytes limit
55 for _ in range(50): # Increase number of entries if needed
56 logger.error(log_message)
57
58 # Ensure the main log file and at least one backup exist
59 assert os.path.exists(log_file), f"Log file {log_file} does not exist."
60 assert os.path.exists(f"{log_file}.1"), f"Backup log file {log_file}.1 does not exist."
```
I know there's a lot there, and the TDD purists among you are probably having an apoplectic fit at the moment (I can hear it now: You can't write all these tests in one shot and be doing TDD right!). Please don't misunderstand me, the design and implementation of these tests took considerable time and didn't happen in a vacuum. Each test was designed individually, with the satisfying function developed and modified after each test. Sometimes, the original tests were inadequate or poorly designed and needed reworking. Other tests made a previous implementation fail a prior test and required rethinking the overall approach to allow both tests to pass. This is the most fun and, honestly, the most challenging part for me. I love puzzles, and trying to figure out how to make all these tests simultaneously true is a giant puzzle. The hardest part was probably making the data collection tests pass, specifically the tests for corrupt data, empty responses, and invalid data types. Finding a way through that maze to ensure all three were covered took the most time.
### **Writing Better Functions Based On Tests**
Remember, if you run PyTest right now, there will be a lot of failures because none of the new tests have complementary implementations to meet the requirements. Now that you can see how the tests are designed and what we're testing for let's take a look at how each one was solved.
Let's start by taking a look at `error_handling.py` where I added some new custom exceptions and updated the logging to handle different error levels:
PYTHON
```
1# IntelliSIEM Copyright 2024, Rob Perry
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15import logging
16from logging.handlers import RotatingFileHandler
17
18# Set up rotating file handler
19handler = RotatingFileHandler('data/error.log', maxBytes=1000000, backupCount=5)
20handler.setLevel(logging.ERROR)
21handler.setFormatter(logging.Formatter('%(asctime)s %(levelname)s: %(message)s'))
22
23# Add the handler to the root logger
24logging.basicConfig(level=logging.ERROR, handlers=[handler])
25
26class APIError(Exception):
27 """Custom exception for API-related errors."""
28 pass
29
30class ConfigError(Exception):
31 """Custom exception for configuration-related errors."""
32 pass
33
34class DataError(Exception):
35 """Custom exception for data processing errors."""
36 pass
37
38def log_error(message, level="ERROR"):
39 """
40 Log error messages to a file with severity level.
41
42 :param message: (str) The error message to log.
43 :param level: (str) The severity level (default is "ERROR").
44 """
45 if level == "WARNING":
46 logging.warning(message)
47 elif level == "INFO":
48 logging.info(message)
49 else:
50 logging.error(message)
```
Along with the new exceptions and better logging, I also added the rotating file handler to keep log files under control. While the test handler was built with a 200KB, the actual implementation uses a 1MB cutoff to trigger rotation.
The next area of focus was the config.py implementation. Here, the focus was predominantly on config file issues such as a missing file, insufficient permissions, or parsing errors.
PYTHON
```
1# IntelliSIEM Copyright 2024, Rob Perry
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15import yaml
16from src.error_handling import log_error, ConfigError
17
18def load_config(config_file='config/config.yaml'):
19 """
20 Load configuration from the specified YAML file.
21
22 :param config_file: (str) The path to the YAML configuration file.
23 :return: dict: Configuration data.
24 :raises ConfigError: If there is an error with the configuration file.
25 """
26 try:
27 with open(config_file, 'r') as file:
28 return yaml.safe_load(file)
29 except yaml.YAMLError as e:
30 log_error(f"YAML parsing error while loading configuration: {e}.")
31 raise ConfigError(f"Configuration file error: {e}.")
32 except FileNotFoundError:
33 log_error(f"Configuration file not found: {config_file}.")
34 raise FileNotFoundError(f"Configuration file not found: {config_file}.")
35 except PermissionError:
36 log_error(f"Permission denied for config file: {config_file}.")
37 raise ConfigError("Permission denied for config file.")
38 except OSError as e:
39 log_error(f"Disk space issue while loading configuration: {e}.")
40 raise ConfigError("Disk space issue while loading configuration.")
```
I anticipate quite a bit more work in this implementation when I circle back to implement stronger security controls around the API keys.
The API client itself took quite a bit of reworking, especially when working through the different exceptions and how to isolate the tests correctly to ensure I was getting valid results:
PYTHON
```
1# IntelliSIEM Copyright 2024, Rob Perry
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15import requests
16from src.error_handling import APIError, log_error
17
18class APIClient:
19 """
20 A client for interacting with various threat intelligence APIs.
21 """
22
23 def __init__(self, base_url, headers):
24 """
25 Initialize the API client with the base URL and headers.
26
27 :param base_url: (str) The base URL of the API.
28 :param headers: (dict) Headers required for API requests.
29 """
30 self.base_url = base_url
31 self.headers = headers
32
33 def get_data(self, endpoint, params=None):
34 """
35 Fetch data from the specified endpoint.
36
37 :param endpoint: (str) The API endpoint to fetch data from.
38 :param params: (dict) Optional query parameters.
39 :return: (dict) JSON response from the API.
40 :raises APIError: If there is an error with the API request.
41 """
42 url = f"{self.base_url}/{endpoint}"
43 try:
44 response = requests.get(url, headers=self.headers, params=params, timeout=10)
45 if response.status_code == 429:
46 retry_after = response.headers.get("Retry-After", "unknown time")
47 log_error(f"Rate limit exceeded. Retry after: {retry_after}.")
48 raise APIError("Rate limit exceeded.")
49 if response.status_code == 401:
50 log_error("Authentication failed: Invalid API key.")
51 raise APIError("Authentication failed.")
52 if response.status_code == 403:
53 log_error("Permission denied: Access to resource is forbidden.")
54 raise APIError("Permission denied.")
55 response.raise_for_status()
56 return response.json()
57 except requests.exceptions.Timeout:
58 log_error("API request timed out.")
59 raise APIError("API request timed out.")
60 except requests.exceptions.ConnectionError:
61 log_error("Connection error occurred.")
62 raise APIError("Connection error occurred.")
63 except requests.exceptions.RequestException as e:
64 log_error(f"API request failed: {e}.")
65 raise APIError(f"API request failed: {e}.")
66 except ValueError:
67 log_error("Invalid JSON response format from API.")
68 raise APIError("Invalid response format.")
```
This particular piece of the implementation should be pretty close to wrapped up now, but I won't feel 100% about it until after some more "live-fire" testing.
The final piece that took a lot of reworking is the data collection function, especially the data validation. Getting the proper order and distinguishing between missing and corrupt data was particularly challenging and effectively required nearly a complete rewrite of the function. The great thing about it, as it is now, is that after I got it to this point, it looks like it'll be a fairly straightforward endeavour to refactor again and make it relatively easy to allow for adding and removing threat feeds based on user preference.
PYTHON
```
1# IntelliSIEM Copyright 2024, Rob Perry
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15from src.api_client import APIClient
16from src.error_handling import APIError, DataError, log_error
17from src.config import load_config
18
19config = load_config()
20
21def fetch_threat_data():
22 """
23 Fetch threat data from various sources and return the combined result.
24
25 :return: (list) A list of threat intelligence data.
26 :raises DataError: If there is an error with data integrity or structure.
27 """
28 try:
29 alien_vault = APIClient("https://otx.alienvault.com/api/v1",
30 {"X-OTX-API-KEY": config['api_keys']['alien_vault']})
31 vt = APIClient("https://www.virustotal.com/api/v3", {"x-apikey": config['api_keys']['virus_total']})
32
33 alien_vault_data = alien_vault.get_data("indicators/export")
34 vt_data = vt.get_data("files", params={"limit": 10})
35
36 # Validate that the data is a list
37 if not isinstance(alien_vault_data, list) or not isinstance(vt_data, list):
38 log_error(f"Invalid data type received: {type(alien_vault_data)} or {type(vt_data)}.")
39 raise DataError("Unexpected data type received.")
40
41 # Define the required keys for validation
42 required_keys = ['id']
43
44 # Validate data and raise DataError if required keys are missing or if data is corrupt
45 valid_data = []
46 for entry in alien_vault_data + vt_data:
47 if not isinstance(entry, dict):
48 log_error(f"Unexpected data type in entry: {entry}")
49 raise DataError("Unexpected data type received.")
50 # Check if the 'id' key is present and its value is of expected type
51 if 'id' in entry and not isinstance(entry['id'], (str, int)):
52 log_error(f"Corrupt data in entry: {entry}. 'id' key has invalid type.")
53 raise DataError("Corrupt data received.")
54 if any(key not in entry for key in required_keys):
55 log_error(f"Missing required data keys in entry: {entry}.")
56 raise DataError("Missing required data keys.")
57 valid_data.append(entry)
58
59 return valid_data
60
61 except APIError as e:
62 log_error(f"Failed to fetch threat data: {e}")
63 return []
64 except (KeyError, TypeError, ValueError) as e:
65 log_error(f"Data error occurred: {e}")
66 raise DataError("Corrupt data received.")
```
I find this particular implementation rather clunky at the moment, but it works. I think implementing a better structure for feed sources will make a significant improvement, but this is going on the TODO list for later.
Implementing robust error handling and log rotation features has significantly improved IntelliSIEM’s resilience. The system can now gracefully handle unexpected data scenarios, and we can track issues effectively through logs without worrying about file sizes. Next, we will turn our attention to beginning the basic UI implementation. Stay tuned as we continue to build out this powerful tool!
---
## Navigation
- Parent: [Blog](https://bsb.perryconsulting.net/blog.md)
- Next: [Implementing TDD in IntelliSIEM: Handling Edge Cases and Errors](https://bsb.perryconsulting.net/blog/strengthening-the-foundation-with-comprehensive-testing-in-intellisiem.md)