def test_resolve_simple_or(): config = [{ "print_name": "Mod_1", "class_name": "mod_1", "vulns": [{ "name": "VULN_1_B", "desc": "Test Vuln 1B", "provides": "", "version": "", "deps": [("B", "=1.2.0"), ("C", "=4.5.0")] }] }, { "print_name": "Mod_2", "class_name": "mod_2", "vulns": [{ "name": "VULN_2_A", "desc": "Test Vuln 2A", "provides": "B", "version": "1.2.0", "deps": [] }] }, { "print_name": "Mod_3", "class_name": "mod_3", "vulns": [{ "name": "VULN_3_A", "desc": "Test Vuln 3A", "provides": "C", "version": "4.5.0", "deps": [] }] }] generate_modules(config) mod_list = module_util.get_module_list() assert "mod_1" in mod_list assert "mod_2" in mod_list assert "mod_3" in mod_list resolving = resolve.resolver() resolving.add_module("mod_1") assert resolving.start_resolve() == True assert resolving.get_install_order() == [ 'mod_3', 'mod_1' ] or resolving.get_install_order() == ['mod_2', 'mod_1'] assert len(resolving.get_install_modules( resolving.get_install_order())) == 2 remove_stubs()
def test_invalid_2(): config = [ { "print_name": "Mod_1", "class_name": "mod_1", "vulns": [ { "name": "VULN_1_A", "desc": "Test Vuln 1A", "provides": "A", "version": "1.0.0", "deps": [ ("B", ">2.0.0") ] } ] }, { "print_name": "Mod_2", "class_name": "mod_2", "vulns": [ { "name": "VULN_2_A", "desc": "Test Vuln 2A", "provides": "B", "version": "1.2.0", "deps": [ ] } ] } ] generate_modules(config) mod_list = module_util.get_module_list() assert "mod_1" in mod_list assert "mod_2" in mod_list resolving = resolve.resolver() resolving.add_module("mod_1") assert resolving.start_resolve() == False resolving = resolve.resolver() resolving.add_module("mod_1") resolving.add_module("mod_2") assert resolving.start_resolve() == False remove_stubs()
def test_resolve_invalid_or(): config = [{ "print_name": "Mod_1", "class_name": "mod_1", "vulns": [{ "name": "VULN_1_A", "desc": "Test Vuln 1A", "provides": "A", "version": "1.0.0", "deps": [("B", ">9.0"), ("C", ">9.0")] }] }, { "print_name": "Mod_2", "class_name": "mod_2", "vulns": [{ "name": "VULN_2_A", "desc": "Test Vuln 2A", "provides": "B", "version": "1.2.0", "deps": [] }] }, { "print_name": "Mod_3", "class_name": "mod_3", "vulns": [{ "name": "VULN_3_A", "desc": "Test Vuln 3A", "provides": "C", "version": "9.1", "deps": [("Z", "=2.0")] }] }] generate_modules(config) mod_list = module_util.get_module_list() assert "mod_1" in mod_list assert "mod_2" in mod_list assert "mod_3" in mod_list resolving = resolve.resolver() resolving.add_module("mod_1") assert resolving.start_resolve() == False remove_stubs()
def test_set_stub_module(): test_obj = module_util.import_module("test_module") test_obj.override_class_name("test2") assert test_obj.get_class_name() == "test2" test_obj.set_name("test2") assert test_obj.name() == "test2" module_util.set_stub_module("test2", test_obj) assert module_util.module_exists("test2") module_list = module_util.get_module_list() assert "test2" in module_list module_util.remove_stub_module("test2") assert not module_util.module_exists("test2") module_list = module_util.get_module_list() assert not "test2" in module_list
def test_set_stub_module(): test_obj = module_util.import_module("test_module") test_obj.override_class_name("test2") assert test_obj.get_class_name() == "test2" test_obj.set_name("test2") assert test_obj.name() == "test2" module_util.set_stub_module("test2", test_obj) assert module_util.module_exists("test2") module_list = module_util.get_module_list() assert "test2" in module_list module_util.remove_stub_module("test2") assert not module_util.module_exists("test2") module_list = module_util.get_module_list() assert not "test2" in module_list
def test_resolve_simple(): mod_1 = module_util.import_module("test_module") vuln1 = mod_1.new_vulnerability("VULN_1_A", "Test Vuln A", "A", "1.0.0") vuln1.add_dependency(("B", ">0.1")) mod_1.add_vulnerability(vuln1) mod_1.override_class_name("mod_1") mod_1.set_name("Mod_1") assert mod_1.get_class_name() == "mod_1" module_util.set_stub_module("mod_1", mod_1) mod_2 = module_util.import_module("test_module") vuln2 = mod_2.new_vulnerability("VULN_2_A", "Test Vuln B", "B", "1.0.0") mod_2.add_vulnerability(vuln2) mod_2.override_class_name("mod_2") mod_2.set_name("Mod_2") assert mod_2.get_class_name() == "mod_2" module_util.set_stub_module("mod_2", mod_2) mod_list = module_util.get_module_list() assert "mod_1" in mod_list assert "mod_2" in mod_list resolving = resolve.resolver() resolving.add_module("mod_1") assert resolving.start_resolve() == True assert resolving.get_install_order() == ['mod_2', 'mod_1'] resolving = resolve.resolver() resolving.add_module("mod_1") resolving.add_module("mod_2") assert resolving.start_resolve() == True assert resolving.get_install_order() == ['mod_2', 'mod_1'] assert resolving.get_install_order() != ['mod_1', 'mod_2'] module_util.remove_stub_module("mod_1") module_util.remove_stub_module("mod_2")
def test_resolve_simple(): mod_1 = module_util.import_module("test_module") vuln1 = mod_1.new_vulnerability("VULN_1_A", "Test Vuln A", "A", "1.0.0") vuln1.add_dependency(("B", ">0.1")) mod_1.add_vulnerability(vuln1) mod_1.override_class_name("mod_1") mod_1.set_name("Mod_1") assert mod_1.get_class_name() == "mod_1" module_util.set_stub_module("mod_1", mod_1) mod_2 = module_util.import_module("test_module") vuln2 = mod_2.new_vulnerability("VULN_2_A", "Test Vuln B", "B", "1.0.0") mod_2.add_vulnerability(vuln2) mod_2.override_class_name("mod_2") mod_2.set_name("Mod_2") assert mod_2.get_class_name() == "mod_2" module_util.set_stub_module("mod_2", mod_2) mod_list = module_util.get_module_list() assert "mod_1" in mod_list assert "mod_2" in mod_list resolving = resolve.resolver() resolving.add_module("mod_1") assert resolving.start_resolve() == True assert resolving.get_install_order() == ['mod_2', 'mod_1'] resolving = resolve.resolver() resolving.add_module("mod_1") resolving.add_module("mod_2") assert resolving.start_resolve() == True assert resolving.get_install_order() == ['mod_2', 'mod_1'] assert resolving.get_install_order() != ['mod_1', 'mod_2'] module_util.remove_stub_module("mod_1") module_util.remove_stub_module("mod_2")
def __load_provides(self, provides_string, restriction_list): choose_list = [] for module in module_util.get_module_list(): current_module = module_util.import_module(module) if current_module.has_provides( provides_string) and current_module.get_class_name( ) not in restriction_list: # If the module we selected is already loaded (but not processed), use it if current_module.get_class_name( ) in self.__name_map and self.__name_map[ current_module.get_class_name()] != None: self.__provides_map[ provides_string] = current_module.get_class_name() return self.__name_map[current_module.get_class_name()] else: choose_list.append(module) if choose_list == 0: return None selected_module_name = ba_random().array_random(choose_list) selected_module = module_util.import_module(selected_module_name) if selected_module == None: return None # Set the provides mapping self.__provides_map[provides_string] = selected_module.get_class_name() self.__insert_module(selected_module.get_class_name(), selected_module) return self.__name_map[selected_module.get_class_name()]
def __load_provides(self, provides_string, restriction_list): choose_list = [] for module in module_util.get_module_list(): current_module = module_util.import_module(module) if current_module.has_provides(provides_string) and current_module.get_class_name() not in restriction_list: # If the module we selected is already loaded (but not processed), use it if current_module.get_class_name() in self.__name_map and self.__name_map[current_module.get_class_name()] != None: self.__provides_map[provides_string] = current_module.get_class_name() return self.__name_map[current_module.get_class_name()] else: choose_list.append(module) if choose_list == 0: return None selected_module_name = ba_random().array_random(choose_list) selected_module = module_util.import_module(selected_module_name) if selected_module == None: return None # Set the provides mapping self.__provides_map[provides_string] = selected_module.get_class_name() self.__insert_module(selected_module.get_class_name(), selected_module) return self.__name_map[selected_module.get_class_name()]
def test_resolve_or(): config = [{ "print_name": "Mod_1", "class_name": "mod_1", "vulns": [{ "name": "VULN_1_A", "desc": "Test Vuln 1A", "provides": "A", "version": "1.0.0", "deps": [("B", "<0.2"), ("C", "<4.5.0")] }, { "name": "VULN_1_B", "desc": "Test Vuln 1B", "provides": "", "version": "", "deps": [("B", "=1.2.0"), ("C", "=4.5.0")] }] }, { "print_name": "Mod_2", "class_name": "mod_2", "vulns": [{ "name": "VULN_2_A", "desc": "Test Vuln 2A", "provides": "B", "version": "1.2.0", "deps": [] }] }, { "print_name": "Mod_3", "class_name": "mod_3", "vulns": [{ "name": "VULN_3_A", "desc": "Test Vuln 3A", "provides": "C", "version": "4.5.0", "deps": [] }] }] generate_modules(config) mod_list = module_util.get_module_list() assert "mod_1" in mod_list assert "mod_2" in mod_list assert "mod_3" in mod_list list1 = ['mod_3', 'mod_1'] list2 = ['mod_2', 'mod_1'] list1_count = 0 list2_count = 0 for i in range(20): #~ resolving = resolve.resolver(debug=True) resolving = resolve.resolver() resolving.add_module("mod_1") assert resolving.start_resolve() == True #~ print("\n\END DEBUG\n\n") install_order = resolving.get_install_order() assert install_order == list1 or install_order == list2 if install_order == list1: list1_count += 1 elif install_order == list2: list2_count += 1 assert len(resolving.get_install_modules( resolving.get_install_order())) == 2 assert list1_count > 1 assert list2_count > 1 remove_stubs()
def __module(self, options): if len(options) == 0: print("Incomplete command") print( "Valid sub-commands: add, remove, info, list, random, force, test" ) else: subcommand = options[0] del options[0] if subcommand == "add": if not len(options) == 1: print("module add: No module name") else: module_name = options[0] if module_util.module_exists(module_name): if not module_name in self.__vars['to_run']['value']: self.__vars['to_run']['value'].append(module_name) print("Module added") else: print("Module is already set to be run") else: print("Module '" + module_name + "' does not exist") elif subcommand == "random": rand_list = [] self.__vars['to_run']['value'] = [] while len(rand_list) == 0: mod_list = module_util.get_module_list() set_level = self.__vars['level']['value'] for module in mod_list: # Skip the test module if module == "test_module": continue mod_obj = module_util.import_module(module) if mod_obj: if set_level == "any": if ba_random().will_do(): rand_list.append(module) else: if mod_obj.has_difficulty( set_level) and ba_random().will_do(): rand_list.append(module) else: print("Could not import module " + module) for module_name in rand_list: self.__vars['to_run']['value'].append(module_name) elif subcommand == "remove": if not len(options) == 1: print("module add: No module set") else: module_name = options[0] if module_util.module_exists(module_name): if module_name in self.__vars['to_run']['value']: self.__vars['to_run']['value'].remove(module_name) print("Module removed") else: print("Module is not set to run") else: print("Module '" + module_name + "' does not exist") elif subcommand == "info": if not len(options) == 1: print("module info: No module name") else: module_name = options[0] if module_util.module_exists(module_name): tmp_module = module_util.import_module(module_name) print(tmp_module.info()) print("Vulnerabilities:") for vuln in tmp_module.full_vulnerability_list(): print(" " + vuln) else: print("Module '" + module_name + "' does not exist") elif subcommand == "list": module_list = module_util.get_module_list() if len(module_list) > 0: for item in module_list: if item == "test_module": continue if item in self.__vars['to_run']['value']: print("\t+ " + item) else: print("\t- " + item) else: print("\nNo modules are set to run!") elif subcommand == "force": if not len(options) == 2: print("Enter a module and vulnerability to force") else: module = options[0] vuln = options[1] if module_util.module_exists(module): if vuln in module_util.import_module( module).full_vulnerability_list(): if not module in self.__vars['force']['value']: self.__vars['force']['value'][module] = [] if not vuln in self.__vars['force']['value'][ module]: self.__vars['force']['value'][module].append( vuln) else: print( "That vulnerability is already being forced" ) else: print("Module '" + module + "' does not have the vulnerability '" + vuln + "'") else: print("Module '" + module + "' does not exist") elif subcommand == "test": if self.__debug_mode == False: print("'module test' can only be used in debug mode") elif not len(options) == 2: print("Enter a module and vulnerability to test") else: module = options[0] vuln = options[1] if module_util.module_exists(module): module_obj = module_util.import_module(module) if vuln in module_obj.full_vulnerability_list(): result = module_obj.test(vuln) if result == False: print("Test failed") else: print("Test succeeded") else: print("Module '" + module + "' does not have the vulnerability '" + vuln + "'") else: print("Module '" + module + "' does not exist") else: print("module: invalid subcommand '" + subcommand + "'")
def test_module_list(): module_list = module_util.get_module_list() assert "test_module" in module_list
def test_resolve_or(): config = [ { "print_name": "Mod_1", "class_name": "mod_1", "vulns": [ { "name": "VULN_1_A", "desc": "Test Vuln 1A", "provides": "A", "version": "1.0.0", "deps": [ ("B", "<0.2"), ("C", "<4.5.0") ] }, { "name": "VULN_1_B", "desc": "Test Vuln 1B", "provides": "", "version": "", "deps": [ ("B", "=1.2.0"), ("C", "=4.5.0") ] } ] }, { "print_name": "Mod_2", "class_name": "mod_2", "vulns": [ { "name": "VULN_2_A", "desc": "Test Vuln 2A", "provides": "B", "version": "1.2.0", "deps": [ ] } ] }, { "print_name": "Mod_3", "class_name": "mod_3", "vulns": [ { "name": "VULN_3_A", "desc": "Test Vuln 3A", "provides": "C", "version": "4.5.0", "deps": [ ] } ] } ] generate_modules(config) mod_list = module_util.get_module_list() assert "mod_1" in mod_list assert "mod_2" in mod_list assert "mod_3" in mod_list list1 = ['mod_3', 'mod_1'] list2 = ['mod_2', 'mod_1'] list1_count = 0 list2_count = 0 for i in range(20): #~ resolving = resolve.resolver(debug=True) resolving = resolve.resolver() resolving.add_module("mod_1") assert resolving.start_resolve() == True #~ print("\n\END DEBUG\n\n") install_order = resolving.get_install_order() assert install_order == list1 or install_order == list2 if install_order == list1: list1_count += 1 elif install_order == list2: list2_count += 1 assert len(resolving.get_install_modules(resolving.get_install_order())) == 2 assert list1_count > 1 assert list2_count > 1 remove_stubs()
def test_resolve_simple_or(): config = [ { "print_name": "Mod_1", "class_name": "mod_1", "vulns": [ { "name": "VULN_1_B", "desc": "Test Vuln 1B", "provides": "", "version": "", "deps": [ ("B", "=1.2.0"), ("C", "=4.5.0") ] } ] }, { "print_name": "Mod_2", "class_name": "mod_2", "vulns": [ { "name": "VULN_2_A", "desc": "Test Vuln 2A", "provides": "B", "version": "1.2.0", "deps": [ ] } ] }, { "print_name": "Mod_3", "class_name": "mod_3", "vulns": [ { "name": "VULN_3_A", "desc": "Test Vuln 3A", "provides": "C", "version": "4.5.0", "deps": [ ] } ] } ] generate_modules(config) mod_list = module_util.get_module_list() assert "mod_1" in mod_list assert "mod_2" in mod_list assert "mod_3" in mod_list resolving = resolve.resolver() resolving.add_module("mod_1") assert resolving.start_resolve() == True assert resolving.get_install_order() == ['mod_3', 'mod_1'] or resolving.get_install_order() == ['mod_2', 'mod_1'] assert len(resolving.get_install_modules(resolving.get_install_order())) == 2 remove_stubs()
def __module(self, options): if len(options) == 0: print("Incomplete command") print("Valid sub-commands: add, remove, info, list, random, force, test") else: subcommand = options[0] del options[0] if subcommand == "add": if not len(options) == 1: print("module add: No module name") else: module_name = options[0] if module_util.module_exists(module_name): if not module_name in self.__vars['to_run']['value']: self.__vars['to_run']['value'].append(module_name) print("Module added") else: print("Module is already set to be run") else: print("Module '" + module_name + "' does not exist") elif subcommand == "random": rand_list = [] self.__vars['to_run']['value'] = [] while len(rand_list) == 0: mod_list = module_util.get_module_list() set_level = self.__vars['level']['value'] for module in mod_list: # Skip the test module if module == "test_module": continue mod_obj = module_util.import_module(module) if mod_obj: if set_level == "any": if ba_random().will_do(): rand_list.append(module) else: if mod_obj.has_difficulty(set_level) and ba_random().will_do(): rand_list.append(module) else: print("Could not import module "+ module) for module_name in rand_list: self.__vars['to_run']['value'].append(module_name) elif subcommand == "remove": if not len(options) == 1: print("module add: No module set") else: module_name = options[0] if module_util.module_exists(module_name): if module_name in self.__vars['to_run']['value']: self.__vars['to_run']['value'].remove(module_name) print("Module removed") else: print("Module is not set to run") else: print("Module '" + module_name + "' does not exist") elif subcommand == "info": if not len(options) == 1: print("module info: No module name") else: module_name = options[0] if module_util.module_exists(module_name): tmp_module = module_util.import_module(module_name) print(tmp_module.info()) print("Vulnerabilities:") for vuln in tmp_module.full_vulnerability_list(): print(" " + vuln) else: print("Module '" + module_name + "' does not exist") elif subcommand == "list": module_list = module_util.get_module_list() if len(module_list) > 0: for item in module_list: if item == "test_module": continue if item in self.__vars['to_run']['value']: print("\t+ " + item) else: print("\t- " + item) else: print("\nNo modules are set to run!") elif subcommand == "force": if not len(options) == 2: print("Enter a module and vulnerability to force") else: module = options[0] vuln = options[1] if module_util.module_exists(module): if vuln in module_util.import_module(module).full_vulnerability_list(): if not module in self.__vars['force']['value']: self.__vars['force']['value'][module] = [] if not vuln in self.__vars['force']['value'][module]: self.__vars['force']['value'][module].append(vuln) else: print("That vulnerability is already being forced") else: print("Module '" + module + "' does not have the vulnerability '" + vuln + "'") else: print("Module '" + module + "' does not exist") elif subcommand == "test": if self.__debug_mode == False: print("'module test' can only be used in debug mode") elif not len(options) == 2: print("Enter a module and vulnerability to test") else: module = options[0] vuln = options[1] if module_util.module_exists(module): module_obj = module_util.import_module(module) if vuln in module_obj.full_vulnerability_list(): result = module_obj.test(vuln) if result == False: print("Test failed") else: print("Test succeeded") else: print("Module '" + module + "' does not have the vulnerability '" + vuln + "'") else: print("Module '" + module + "' does not exist") else: print("module: invalid subcommand '" + subcommand + "'")
def test_module_list(): module_list = module_util.get_module_list() assert "test_module" in module_list