diff --git a/ScoutSuite/providers/base/provider.py b/ScoutSuite/providers/base/provider.py index 64c4bbd05..7aae732bd 100755 --- a/ScoutSuite/providers/base/provider.py +++ b/ScoutSuite/providers/base/provider.py @@ -192,19 +192,29 @@ def recursive_get_count(self, resource, resources): def manage_object(self, object, attr, init, callback=None): """ - This is a quick-fix copy of Opinel's manage_dictionary in order to support the new ScoutSuite object which isn't - a dict + Initialize an attribute in an object or dictionary if it doesn't exist. + This is a quick-fix copy of Opinel's manage_dictionary modified to support both + dictionaries and objects while avoiding infinite recursion. + + :param object: The dictionary or object to modify + :param attr: The attribute name + :param init: The initial value + :param callback: Optional callback to execute after setting the value + :return: The modified object """ if type(object) == dict: if not str(attr) in object: object[str(attr)] = init - self.manage_object(object, attr, init) else: if not hasattr(object, attr): setattr(object, attr, init) - self.manage_object(object, attr, init) + if callback: - callback(getattr(object, attr)) + if type(object) == dict: + callback(object[str(attr)]) + else: + callback(getattr(object, attr)) + return object def _process_metadata_callbacks(self): diff --git a/tests/providers/base/test_provider.py b/tests/providers/base/test_provider.py new file mode 100644 index 000000000..9ac69731c --- /dev/null +++ b/tests/providers/base/test_provider.py @@ -0,0 +1,79 @@ +import unittest +from unittest.mock import patch +from ScoutSuite.providers.base.provider import BaseProvider + +class MockBaseProvider(BaseProvider): + """Mock class to test BaseProvider without requiring metadata""" + def __init__(self): + self.metadata = {} + self.last_run = None + self.services = {} + + def _load_metadata(self): + """Override to prevent metadata loading""" + pass + +class TestBaseProvider(unittest.TestCase): + def setUp(self): + """Set up a BaseProvider instance for testing""" + self.provider = MockBaseProvider() + + def test_manage_object_dict(self): + """Test manage_object with dictionary input""" + # Test adding new key to dictionary + test_dict = {} + result = self.provider.manage_object(test_dict, 'new_key', 'test_value') + self.assertEqual(result['new_key'], 'test_value') + + # Test existing key in dictionary + test_dict = {'existing_key': 'old_value'} + result = self.provider.manage_object(test_dict, 'existing_key', 'new_value') + self.assertEqual(result['existing_key'], 'old_value') # Should not change existing value + + def test_manage_object_class(self): + """Test manage_object with class object input""" + class TestClass: + pass + + # Test adding new attribute + test_obj = TestClass() + result = self.provider.manage_object(test_obj, 'new_attr', 'test_value') + self.assertEqual(result.new_attr, 'test_value') + + # Test existing attribute + test_obj = TestClass() + setattr(test_obj, 'existing_attr', 'old_value') + result = self.provider.manage_object(test_obj, 'existing_attr', 'new_value') + self.assertEqual(result.existing_attr, 'old_value') # Should not change existing value + + def test_manage_object_callback(self): + """Test manage_object with callback function""" + callback_called = {'value': None} + + def test_callback(value): + callback_called['value'] = value + + # Test callback with dictionary + test_dict = {} + self.provider.manage_object(test_dict, 'callback_key', 'callback_value', callback=test_callback) + self.assertEqual(callback_called['value'], 'callback_value') + + # Test callback with object + class TestClass: + pass + test_obj = TestClass() + callback_called['value'] = None # Reset callback value + self.provider.manage_object(test_obj, 'callback_attr', 'callback_value', callback=test_callback) + self.assertEqual(callback_called['value'], 'callback_value') + + def test_manage_object_no_infinite_loop(self): + """Test that manage_object doesn't cause infinite recursion""" + # This test would have failed with the original implementation + test_dict = {} + try: + self.provider.manage_object(test_dict, 'test_key', 'test_value') + except RecursionError: + self.fail("manage_object caused infinite recursion") + +if __name__ == '__main__': + unittest.main()