From 62e46dec7bdc7316e570fdc237740b167bd73fa1 Mon Sep 17 00:00:00 2001 From: benyissa Date: Fri, 17 Nov 2023 12:00:44 +0100 Subject: [PATCH 1/6] implement process --- agent/asteroid_agent.py | 22 +++++++++++++- tests/asteroid_agent_test.py | 20 +++++++++++-- tests/conftest.py | 58 ++++++++++++++++++++++++++++++++++++ tests/test-requirement.txt | 1 + 4 files changed, 97 insertions(+), 4 deletions(-) diff --git a/agent/asteroid_agent.py b/agent/asteroid_agent.py index 70cd9fc1..e6cd0db0 100644 --- a/agent/asteroid_agent.py +++ b/agent/asteroid_agent.py @@ -8,6 +8,10 @@ from ostorlab.agent.mixins import agent_report_vulnerability_mixin from ostorlab.agent.message import message as m +from agent import utils +from agent import exploits_registry +from agent import definitions + logging.basicConfig( format="%(message)s", datefmt="[%X]", @@ -30,7 +34,23 @@ def process(self, message: m.Message) -> None: message: message containing the asset to scan. """ - # TODO (benyissa): implement agent logic here. + targets = utils.prepare_targets(message) + exploits_list: list[ + definitions.Exploit + ] = exploits_registry.ExploitsRegistry.values() + for target in targets: + for exploit in exploits_list: + if exploit.accept(target) is False: + continue + vulnz = exploit.check(target) + for vulnerability in vulnz: + self.report_vulnerability( + entry=vulnerability.entry, + risk_rating=vulnerability.risk_rating, + vulnerability_location=vulnerability.vulnerability_location, + dna=vulnerability.dna, + technical_detail=vulnerability.technical_detail, + ) if __name__ == "__main__": diff --git a/tests/asteroid_agent_test.py b/tests/asteroid_agent_test.py index 15c64c39..9a2723ad 100644 --- a/tests/asteroid_agent_test.py +++ b/tests/asteroid_agent_test.py @@ -1,6 +1,20 @@ """Unit tests for AsteroidAgent.""" +from typing import Type +from agent import asteroid_agent +from agent import definitions +from ostorlab.agent.message import message as m -def testAgent() -> None: - """Fake test.""" - assert True + +def testAsteroidAgent_whenNoExploitYet_doesNotEmitsVulnerabilityReport( + asteroid_agent_instance: asteroid_agent.AsteroidAgent, + exploit_instance_with_report: Type[definitions.Exploit], + agent_mock: list[m.Message], + scan_message_domain_name: m.Message, +) -> None: + """Unit test for agent AsteroidAgent exploits check. case Exploit emits vulnerability report""" + + asteroid_agent_instance.process(scan_message_domain_name) + + assert len(agent_mock) == 1 + assert agent_mock[0].selector == "v3.report.vulnerability" diff --git a/tests/conftest.py b/tests/conftest.py index f8e95728..29254546 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,8 +1,16 @@ """Pytest fixtures for agent Asteroid""" +import pathlib +import random from typing import Type import pytest +from ostorlab.agent import definitions as agent_definitions from ostorlab.agent.message import message +from ostorlab.runtimes import definitions as runtime_definitions +from ostorlab.agent.mixins import agent_report_vulnerability_mixin as vuln_mixin +from ostorlab.agent.kb import kb +from agent import asteroid_agent +from agent import exploits_registry from agent import definitions @@ -20,6 +28,40 @@ def check(self, target: definitions.Target) -> list[definitions.Vulnerability]: return TestExploit +@pytest.fixture() +def exploit_instance_with_report() -> Type[definitions.Exploit]: + @exploits_registry.register + class TestExploit(definitions.Exploit): + """test class Exploit.""" + + def accept(self, target: definitions.Target) -> bool: + return True + + def check(self, target: definitions.Target) -> list[definitions.Vulnerability]: + return [ + definitions.Vulnerability( + technical_detail="test", + entry=kb.Entry( + title="test", + risk_rating="INFO", + short_description="test purposes", + description="test purposes", + recommendation="", + references={}, + security_issue=False, + privacy_issue=False, + has_public_exploit=False, + targeted_by_malware=False, + targeted_by_ransomware=False, + targeted_by_nation_state=False, + ), + risk_rating=vuln_mixin.RiskRating.HIGH, + ) + ] + + return TestExploit + + @pytest.fixture() def scan_message_domain_name() -> message.Message: """Creates a message of type v3.asset.domain_name.service to be used by the agent for testing purposes.""" @@ -54,3 +96,19 @@ def scan_message_ipv4() -> message.Message: selector = "v3.asset.ip.v4" msg_data = {"host": "192.168.1.17", "mask": "32", "version": 4} return message.Message.from_data(selector, data=msg_data) + + +@pytest.fixture() +def asteroid_agent_instance() -> asteroid_agent.AsteroidAgent: + with (pathlib.Path(__file__).parent.parent / "ostorlab.yaml").open() as yaml_o: + definition = agent_definitions.AgentDefinition.from_yaml(yaml_o) + settings = runtime_definitions.AgentSettings( + key="agent/ostorlab/asteroid", + bus_url="NA", + bus_exchange_topic="NA", + args=[], + healthcheck_port=random.randint(5000, 6000), + redis_url="redis://guest:guest@localhost:6379", + ) + + return asteroid_agent.AsteroidAgent(definition, settings) diff --git a/tests/test-requirement.txt b/tests/test-requirement.txt index a85bde23..50bc9271 100644 --- a/tests/test-requirement.txt +++ b/tests/test-requirement.txt @@ -2,4 +2,5 @@ pylint pytest black mypy +pytest_mock typing-extensions From 2020935836a92251db3fcfcdf062b3b56e96242b Mon Sep 17 00:00:00 2001 From: benyissa Date: Fri, 17 Nov 2023 15:05:07 +0100 Subject: [PATCH 2/6] update test name --- tests/asteroid_agent_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/asteroid_agent_test.py b/tests/asteroid_agent_test.py index 9a2723ad..dbf93609 100644 --- a/tests/asteroid_agent_test.py +++ b/tests/asteroid_agent_test.py @@ -6,7 +6,7 @@ from ostorlab.agent.message import message as m -def testAsteroidAgent_whenNoExploitYet_doesNotEmitsVulnerabilityReport( +def testAsteroidAgent_whenExploitCheckDetectVulnz_EmitsVulnerabilityReport( asteroid_agent_instance: asteroid_agent.AsteroidAgent, exploit_instance_with_report: Type[definitions.Exploit], agent_mock: list[m.Message], From 3e3e27f49e1cc0b477ecd394f5df06fe0f8004ba Mon Sep 17 00:00:00 2001 From: benyissa Date: Mon, 20 Nov 2023 13:11:00 +0100 Subject: [PATCH 3/6] use ThreadPoolExecutor --- agent/asteroid_agent.py | 36 ++++++++++++++++++++++++++---------- tests/asteroid_agent_test.py | 2 +- 2 files changed, 27 insertions(+), 11 deletions(-) diff --git a/agent/asteroid_agent.py b/agent/asteroid_agent.py index e6cd0db0..c18af3fe 100644 --- a/agent/asteroid_agent.py +++ b/agent/asteroid_agent.py @@ -3,9 +3,13 @@ `v3.report.vulnerability` with a technical report.""" import logging from rich import logging as rich_logging +import concurrent.futures + from ostorlab.agent import agent +from ostorlab.agent import definitions as agent_definitions from ostorlab.agent.mixins import agent_report_vulnerability_mixin +from ostorlab.runtimes import definitions as runtime_definitions from ostorlab.agent.message import message as m from agent import utils @@ -25,6 +29,18 @@ class AsteroidAgent(agent.Agent, agent_report_vulnerability_mixin.AgentReportVulnMixin): """Asteroid Agent is designed to identify known exploitable vulnerabilities in a remote system.""" + def __init__( + self, + agent_definition: agent_definitions.AgentDefinition, + agent_settings: runtime_definitions.AgentSettings, + ) -> None: + """Initialize The Agent instance.""" + + super().__init__(agent_definition, agent_settings) + self.exploits: list[ + definitions.Exploit + ] = exploits_registry.ExploitsRegistry.values() + def process(self, message: m.Message) -> None: """Process messages of type `v3.asset.ip.[v4,v6]` or `v3.asset.[domain_name,link]` and performs a network scan. Once the scan is completed, it emits messages of type @@ -33,17 +49,17 @@ def process(self, message: m.Message) -> None: Args: message: message containing the asset to scan. """ - targets = utils.prepare_targets(message) - exploits_list: list[ - definitions.Exploit - ] = exploits_registry.ExploitsRegistry.values() - for target in targets: - for exploit in exploits_list: - if exploit.accept(target) is False: - continue - vulnz = exploit.check(target) - for vulnerability in vulnz: + with concurrent.futures.ThreadPoolExecutor() as executor: + for exploit in self.exploits: + for target in targets: + exploit_vulnz = [] + if exploit.accept(target) is False: + continue + exploit_vulnz.append(executor.submit(exploit.check, target)) + + for target_vulnz in exploit_vulnz: + for vulnerability in target_vulnz.result(): self.report_vulnerability( entry=vulnerability.entry, risk_rating=vulnerability.risk_rating, diff --git a/tests/asteroid_agent_test.py b/tests/asteroid_agent_test.py index dbf93609..8dbec291 100644 --- a/tests/asteroid_agent_test.py +++ b/tests/asteroid_agent_test.py @@ -7,8 +7,8 @@ def testAsteroidAgent_whenExploitCheckDetectVulnz_EmitsVulnerabilityReport( - asteroid_agent_instance: asteroid_agent.AsteroidAgent, exploit_instance_with_report: Type[definitions.Exploit], + asteroid_agent_instance: asteroid_agent.AsteroidAgent, agent_mock: list[m.Message], scan_message_domain_name: m.Message, ) -> None: From ef0c6fbe8d8201e294e91752631defb4149eb86e Mon Sep 17 00:00:00 2001 From: benyissa Date: Mon, 20 Nov 2023 13:20:34 +0100 Subject: [PATCH 4/6] fix after merge --- agent/asteroid_agent.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/agent/asteroid_agent.py b/agent/asteroid_agent.py index c18af3fe..90a88610 100644 --- a/agent/asteroid_agent.py +++ b/agent/asteroid_agent.py @@ -12,7 +12,7 @@ from ostorlab.runtimes import definitions as runtime_definitions from ostorlab.agent.message import message as m -from agent import utils +from agent import targets_preparer from agent import exploits_registry from agent import definitions @@ -49,7 +49,7 @@ def process(self, message: m.Message) -> None: Args: message: message containing the asset to scan. """ - targets = utils.prepare_targets(message) + targets = targets_preparer.prepare_targets(message) with concurrent.futures.ThreadPoolExecutor() as executor: for exploit in self.exploits: for target in targets: From 58cb9ca7ca29c9ebabb648e51af1bcb7a3beb26f Mon Sep 17 00:00:00 2001 From: benyissa Date: Mon, 20 Nov 2023 15:44:02 +0100 Subject: [PATCH 5/6] fix comments --- agent/asteroid_agent.py | 23 +++++++++++++++-------- 1 file changed, 15 insertions(+), 8 deletions(-) diff --git a/agent/asteroid_agent.py b/agent/asteroid_agent.py index 90a88610..4b814455 100644 --- a/agent/asteroid_agent.py +++ b/agent/asteroid_agent.py @@ -3,7 +3,7 @@ `v3.report.vulnerability` with a technical report.""" import logging from rich import logging as rich_logging -import concurrent.futures +from concurrent import futures from ostorlab.agent import agent @@ -26,6 +26,14 @@ logger = logging.getLogger(__name__) +def _check_target( + exploit: definitions.Exploit, target: definitions.Target +) -> list[definitions.Vulnerability]: + if exploit.accept(target) is False: + return [] + return exploit.check(target) + + class AsteroidAgent(agent.Agent, agent_report_vulnerability_mixin.AgentReportVulnMixin): """Asteroid Agent is designed to identify known exploitable vulnerabilities in a remote system.""" @@ -50,15 +58,14 @@ def process(self, message: m.Message) -> None: message: message containing the asset to scan. """ targets = targets_preparer.prepare_targets(message) - with concurrent.futures.ThreadPoolExecutor() as executor: + with futures.ThreadPoolExecutor() as executor: for exploit in self.exploits: - for target in targets: - exploit_vulnz = [] - if exploit.accept(target) is False: - continue - exploit_vulnz.append(executor.submit(exploit.check, target)) + futures_exploit_check = [ + executor.submit(_check_target, exploit, target) + for target in targets + ] - for target_vulnz in exploit_vulnz: + for target_vulnz in futures.as_completed(futures_exploit_check): for vulnerability in target_vulnz.result(): self.report_vulnerability( entry=vulnerability.entry, From e9ca1c79b015ed1f14ff21686518a8a64b27f5df Mon Sep 17 00:00:00 2001 From: benyissa Date: Mon, 20 Nov 2023 15:58:45 +0100 Subject: [PATCH 6/6] fix comments --- agent/asteroid_agent.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/agent/asteroid_agent.py b/agent/asteroid_agent.py index 4b814455..447ae3c8 100644 --- a/agent/asteroid_agent.py +++ b/agent/asteroid_agent.py @@ -5,7 +5,6 @@ from rich import logging as rich_logging from concurrent import futures - from ostorlab.agent import agent from ostorlab.agent import definitions as agent_definitions from ostorlab.agent.mixins import agent_report_vulnerability_mixin @@ -59,13 +58,13 @@ def process(self, message: m.Message) -> None: """ targets = targets_preparer.prepare_targets(message) with futures.ThreadPoolExecutor() as executor: - for exploit in self.exploits: - futures_exploit_check = [ - executor.submit(_check_target, exploit, target) - for target in targets - ] + targets_checks = [ + executor.submit(_check_target, exploit, target) + for target in targets + for exploit in self.exploits + ] - for target_vulnz in futures.as_completed(futures_exploit_check): + for target_vulnz in futures.as_completed(targets_checks): for vulnerability in target_vulnz.result(): self.report_vulnerability( entry=vulnerability.entry,