diff --git a/src/core/src/service_interfaces/StatusHandler.py b/src/core/src/service_interfaces/StatusHandler.py index d25233cd..dbf8ae23 100644 --- a/src/core/src/service_interfaces/StatusHandler.py +++ b/src/core/src/service_interfaces/StatusHandler.py @@ -20,6 +20,7 @@ import shutil import time from core.src.bootstrap.Constants import Constants +from collections import OrderedDict class StatusHandler(object): @@ -53,6 +54,7 @@ def __init__(self, env_layer, execution_config, composite_logger, telemetry_writ self.__assessment_packages = [] self.__assessment_errors = [] self.__assessment_total_error_count = 0 # All errors during assess, includes errors not in error objects due to size limit + self.__assessment_packages_map = OrderedDict() # Internal in-memory representation of Patch Metadata for HealthStore self.__metadata_for_healthstore_substatus_json = None @@ -101,18 +103,21 @@ def reset_assessment_data(self): self.__assessment_packages = [] self.__assessment_errors = [] self.__assessment_total_error_count = 0 + self.__assessment_packages_map = OrderedDict() def set_package_assessment_status(self, package_names, package_versions, classification="Other", status="Available"): """ Externally available method to set assessment status for one or more packages of the **SAME classification and status** """ self.composite_logger.log_debug("Setting package assessment status in bulk. [Count={0}]".format(str(len(package_names)))) + for package_name, package_version in zip(package_names, package_versions): patch_already_saved = False patch_id = self.__get_patch_id(package_name, package_version) - for i in range(0, len(self.__assessment_packages)): - if patch_id == self.__assessment_packages[i]['patchId']: - patch_already_saved = True - self.__assessment_packages[i]['classifications'] = [classification] - # self.__assessment_packages[i]['patchState'] = status + + # Match patch_id in map and update existing patch's classification i.e from other -> security + if len(self.__assessment_packages_map) > 0 and patch_id in self.__assessment_packages_map: + self.__assessment_packages_map.setdefault(patch_id, {})['classifications'] = [classification] + # self.__assessment_packages_map.setdefault(patch_id, {})['patchState'] = status + patch_already_saved = True if patch_already_saved is False: record = { @@ -122,8 +127,10 @@ def set_package_assessment_status(self, package_names, package_versions, classif "classifications": [classification] # "patchState": str(status) # Allows for capturing 'Installed' packages in addition to 'Available', when commented out, if spec changes } - self.__assessment_packages.append(record) + # Add new patch to map + self.__assessment_packages_map[patch_id] = record + self.__assessment_packages = list(self.__assessment_packages_map.values()) self.__assessment_packages = self.sort_packages_by_classification_and_state(self.__assessment_packages) self.set_assessment_substatus_json() @@ -540,6 +547,7 @@ def load_status_file_components(self, initial_load=False): self.__assessment_summary_json = None self.__assessment_packages = [] self.__assessment_errors = [] + self.__assessment_packages_map = OrderedDict() self.__metadata_for_healthstore_substatus_json = None self.__metadata_for_healthstore_summary_json = None @@ -601,7 +609,8 @@ def load_status_file_components(self, initial_load=False): if name == Constants.PATCH_ASSESSMENT_SUMMARY: # if it exists, it must be to spec, or an exception will get thrown message = status_file_data['status']['substatus'][i]['formattedMessage']['message'] self.__assessment_summary_json = json.loads(message) - self.__assessment_packages = self.__assessment_summary_json['patches'] + self.__assessment_packages_map = OrderedDict((package["patchId"], package) for package in self.__assessment_summary_json['patches']) + self.__assessment_packages = list(self.__assessment_packages_map.values()) errors = self.__assessment_summary_json['errors'] if errors is not None and errors['details'] is not None: self.__assessment_errors = errors['details'] diff --git a/src/core/tests/Test_StatusHandler.py b/src/core/tests/Test_StatusHandler.py index 657adb1d..44f2be2d 100644 --- a/src/core/tests/Test_StatusHandler.py +++ b/src/core/tests/Test_StatusHandler.py @@ -396,6 +396,37 @@ def test_sort_packages_by_classification_and_state(self): self.assertEqual(installation_patches_sorted[12]["name"], "test-package-2") # | Other | Excluded | self.assertEqual(installation_patches_sorted[13]["name"], "test-package-1") # | Other | NotSelected | + def test_assessment_packages_map(self): + patch_count_for_test = 5 + expected_patch_id = 'python-samba0_2:4.4.5+dfsg-2ubuntu5.4_Ubuntu_16.04' + + status_handler = StatusHandler(self.runtime.env_layer, self.runtime.execution_config, self.runtime.composite_logger, self.runtime.telemetry_writer, self.runtime.vm_cloud_type) + self.runtime.execution_config.operation = Constants.ASSESSMENT + self.runtime.status_handler.set_current_operation(Constants.ASSESSMENT) + + test_packages, test_package_versions = self.__set_up_packages_func(patch_count_for_test) + status_handler.set_package_assessment_status(test_packages, test_package_versions, 'Critical') + + with self.runtime.env_layer.file_system.open(self.runtime.execution_config.status_file_path, 'r') as file_handle: + substatus_file_data = json.load(file_handle)[0]["status"]["substatus"][0] + + self.assertTrue(substatus_file_data["name"] == Constants.PATCH_ASSESSMENT_SUMMARY) + formatted_message = json.loads(substatus_file_data['formattedMessage']['message']) + self.assertEqual(len(formatted_message['patches']), patch_count_for_test) + self.assertEqual(formatted_message['patches'][0]['classifications'], ['Critical']) + self.assertEqual(formatted_message['patches'][0]['name'], 'python-samba0') + self.assertEqual(formatted_message['patches'][0]['patchId'], expected_patch_id) + + # Setup functions to populate packages and versions for truncation + def __set_up_packages_func(self, val): + test_packages = [] + test_package_versions = [] + + for i in range(0, val): + test_packages.append('python-samba' + str(i)) + test_package_versions.append('2:4.4.5+dfsg-2ubuntu5.4') + + return test_packages, test_package_versions if __name__ == '__main__': unittest.main()