Merge pull request #9 from Ostorlab/implement_process

Feat / Implement process

agent/asteroid_agent.py

@@ -3,11 +3,18 @@
 `v3.report.vulnerability` with a technical report."""
 import logging
 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
+from ostorlab.runtimes import definitions as runtime_definitions
 from ostorlab.agent.message import message as m
 
+from agent import targets_preparer
+from agent import exploits_registry
+from agent import definitions
+
 logging.basicConfig(
     format="%(message)s",
     datefmt="[%X]",
@@ -18,9 +25,29 @@
 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."""
 
+    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
@@ -29,8 +56,23 @@ def process(self, message: m.Message) -> None:
         Args:
             message: message containing the asset to scan.
         """
+        targets = targets_preparer.prepare_targets(message)
+        with futures.ThreadPoolExecutor() as executor:
+            targets_checks = [
+                executor.submit(_check_target, exploit, target)
+                for target in targets
+                for exploit in self.exploits
+            ]
 
-        # TODO (benyissa): implement agent logic here.
+            for target_vulnz in futures.as_completed(targets_checks):
+                for vulnerability in target_vulnz.result():
+                    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__":

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_whenExploitCheckDetectVulnz_EmitsVulnerabilityReport(
+    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:
+    """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"

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)

tests/test-requirement.txt

@@ -2,6 +2,7 @@ pylint
 pytest
 black
 mypy
+pytest_mock
 typing-extensions
 requests_mock
 types-requests