p17860254
/
monkey
forked from p15670423/monkey
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Agent: Accept IAgentEventQueue in Puppet constructor

This commit is contained in:
Ilija Lazoroski 2022-09-28 11:13:05 +02:00 committed by Shreya Malviya
parent ec56b15219
commit 0357d43d33
2 changed files with 8 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
monkey/infection_monkey/puppet/puppet.py
View File

@ -4,6 +4,7 @@ from typing import Dict, Iterable, Sequence
from common.common_consts.timeouts import CONNECTION_TIMEOUT from common.common_consts.timeouts import CONNECTION_TIMEOUT
from common.credentials import Credentials from common.credentials import Credentials
from common.event_queue import IAgentEventQueue
from common.types import PingScanData from common.types import PingScanData
from infection_monkey import network_scanning from infection_monkey import network_scanning
from infection_monkey.i_puppet import ( from infection_monkey.i_puppet import (
@ -24,8 +25,9 @@ logger = logging.getLogger()
class Puppet(IPuppet): class Puppet(IPuppet):
def __init__(self) -> None: def __init__(self, agent_event_queue: IAgentEventQueue) -> None:
self._plugin_registry = PluginRegistry() self._plugin_registry = PluginRegistry()
self._agent_event_queue = agent_event_queue
def load_plugin(self, plugin_name: str, plugin: object, plugin_type: PluginType) -> None: def load_plugin(self, plugin_name: str, plugin: object, plugin_type: PluginType) -> None:
self._plugin_registry.load_plugin(plugin_name, plugin, plugin_type) self._plugin_registry.load_plugin(plugin_name, plugin, plugin_type)
@ -41,7 +43,7 @@ class Puppet(IPuppet):
return pba.run(options) return pba.run(options)
def ping(self, host: str, timeout: float = CONNECTION_TIMEOUT) -> PingScanData: def ping(self, host: str, timeout: float = CONNECTION_TIMEOUT) -> PingScanData:
return network_scanning.ping(host, timeout) return network_scanning.ping(host, timeout, self._agent_event_queue)
def scan_tcp_ports( def scan_tcp_ports(
self, host: str, ports: Sequence[int], timeout: float = CONNECTION_TIMEOUT self, host: str, ports: Sequence[int], timeout: float = CONNECTION_TIMEOUT

7
monkey/tests/unit_tests/infection_monkey/puppet/test_puppet.py
View File

@ -1,13 +1,14 @@
import threading import threading
from unittest.mock import MagicMock from unittest.mock import MagicMock
from common.event_queue import IAgentEventQueue
from common.types import PingScanData from common.types import PingScanData
from infection_monkey.i_puppet import PluginType from infection_monkey.i_puppet import PluginType
from infection_monkey.puppet.puppet import EMPTY_FINGERPRINT, Puppet from infection_monkey.puppet.puppet import EMPTY_FINGERPRINT, Puppet
def test_puppet_run_payload_success(): def test_puppet_run_payload_success():
p = Puppet() p = Puppet(agent_event_queue=MagicMock(spec=IAgentEventQueue))
payload = MagicMock() payload = MagicMock()
payload_name = "PayloadOne" payload_name = "PayloadOne"
@ -19,7 +20,7 @@ def test_puppet_run_payload_success():
def test_puppet_run_multiple_payloads(): def test_puppet_run_multiple_payloads():
p = Puppet() p = Puppet(agent_event_queue=MagicMock(spec=IAgentEventQueue))
payload_1 = MagicMock() payload_1 = MagicMock()
payload1_name = "PayloadOne" payload1_name = "PayloadOne"
@ -45,6 +46,6 @@ def test_puppet_run_multiple_payloads():
def test_fingerprint_exception_handling(monkeypatch): def test_fingerprint_exception_handling(monkeypatch):
p = Puppet() p = Puppet(agent_event_queue=MagicMock(spec=IAgentEventQueue))
p._plugin_registry.get_plugin = MagicMock(side_effect=Exception) p._plugin_registry.get_plugin = MagicMock(side_effect=Exception)
assert p.fingerprint("", "", PingScanData("windows", False), {}, {}) == EMPTY_FINGERPRINT assert p.fingerprint("", "", PingScanData("windows", False), {}, {}) == EMPTY_FINGERPRINT