Skip to content

Commit b058ca5

Browse files
mssalvatoremssalvatore
committed
Merge branch '3167-load-credential-collector-plugins' into develop
Issue #3167 PR #3254
2 parents 3f683f4 + 675d33d commit b058ca5

8 files changed

+226
-122
lines changed

‎monkey/infection_monkey/monkey.py

+27-8
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
from typing import Optional, Sequence, Tuple
1515

1616
from pubsub.core import Publisher
17-
from serpentarium import PluginLoader
17+
from serpentarium import PluginLoader, PluginThreadName
1818
from serpentarium.logging import configure_child_process_logger
1919

2020
from common import HARD_CODED_EXPLOITER_MANIFESTS, OperatingSystem
@@ -82,6 +82,11 @@
8282
from infection_monkey.network_scanning.smb_fingerprinter import SMBFingerprinter
8383
from infection_monkey.network_scanning.ssh_fingerprinter import SSHFingerprinter
8484
from infection_monkey.payload.ransomware.ransomware_payload import RansomwarePayload
85+
from infection_monkey.plugin.credentials_collector_plugin_factory import (
86+
CredentialsCollectorPluginFactory,
87+
)
88+
from infection_monkey.plugin.exploiter_plugin_factory import ExploiterPluginFactory
89+
from infection_monkey.plugin.multiprocessing_plugin_wrapper import MultiprocessingPluginWrapper
8590
from infection_monkey.propagation_credentials_repository import (
8691
AggregatingPropagationCredentialsRepository,
8792
PropagationCredentialsRepository,
@@ -386,17 +391,31 @@ def _build_puppet(self, operating_system: OperatingSystem) -> IPuppet:
386391
self._plugin_dir, partial(configure_child_process_logger, self._ipc_logger_queue)
387392
)
388393
otp_provider = IslandAPIAgentOTPProvider(self._island_api_client)
394+
create_plugin = partial(
395+
MultiprocessingPluginWrapper,
396+
plugin_loader=plugin_loader,
397+
reset_modules_cache=False,
398+
main_thread_name=PluginThreadName.CALLING_THREAD,
399+
)
400+
plugin_factories = {
401+
AgentPluginType.CREDENTIALS_COLLECTOR: CredentialsCollectorPluginFactory(
402+
self._agent_id, self._agent_event_publisher, create_plugin
403+
),
404+
AgentPluginType.EXPLOITER: ExploiterPluginFactory(
405+
self._agent_id,
406+
agent_binary_repository,
407+
self._agent_event_publisher,
408+
self._propagation_credentials_repository,
409+
self._tcp_port_selector,
410+
otp_provider,
411+
create_plugin,
412+
),
413+
}
389414
plugin_registry = PluginRegistry(
390415
operating_system,
391416
self._island_api_client,
392417
plugin_source_extractor,
393-
plugin_loader,
394-
agent_binary_repository,
395-
self._agent_event_publisher,
396-
self._propagation_credentials_repository,
397-
self._tcp_port_selector,
398-
otp_provider,
399-
self._agent_id,
418+
plugin_factories,
400419
)
401420
plugin_compatability_verifier = PluginCompatabilityVerifier(
402421
self._island_api_client, HARD_CODED_EXPLOITER_MANIFESTS

‎monkey/infection_monkey/plugin/credentials_collector_plugin_factory.py

+27
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
from typing import Callable
2+
3+
from serpentarium import SingleUsePlugin
4+
5+
from common.event_queue import IAgentEventPublisher
6+
from common.types import AgentID
7+
8+
from .i_plugin_factory import IPluginFactory
9+
10+
11+
class CredentialsCollectorPluginFactory(IPluginFactory):
12+
def __init__(
13+
self,
14+
agent_id: AgentID,
15+
agent_event_publisher: IAgentEventPublisher,
16+
create_plugin: Callable[..., SingleUsePlugin],
17+
):
18+
self._agent_id = agent_id
19+
self._agent_event_publisher = agent_event_publisher
20+
self._create_plugin = create_plugin
21+
22+
def create(self, plugin_name: str) -> SingleUsePlugin:
23+
return self._create_plugin(
24+
plugin_name=plugin_name,
25+
agent_id=self._agent_id,
26+
agent_event_publisher=self._agent_event_publisher,
27+
)

‎monkey/infection_monkey/plugin/exploiter_plugin_factory.py

+42
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
from typing import Callable
2+
3+
from serpentarium import SingleUsePlugin
4+
5+
from common.event_queue import IAgentEventPublisher
6+
from common.types import AgentID
7+
from infection_monkey.exploit import IAgentBinaryRepository, IAgentOTPProvider
8+
from infection_monkey.network import TCPPortSelector
9+
from infection_monkey.propagation_credentials_repository import IPropagationCredentialsRepository
10+
11+
from .i_plugin_factory import IPluginFactory
12+
13+
14+
class ExploiterPluginFactory(IPluginFactory):
15+
def __init__(
16+
self,
17+
agent_id: AgentID,
18+
agent_binary_repository: IAgentBinaryRepository,
19+
agent_event_publisher: IAgentEventPublisher,
20+
propagation_credentials_repository: IPropagationCredentialsRepository,
21+
tcp_port_selector: TCPPortSelector,
22+
otp_provider: IAgentOTPProvider,
23+
create_plugin: Callable[..., SingleUsePlugin],
24+
):
25+
self._agent_id = agent_id
26+
self._agent_binary_repository = agent_binary_repository
27+
self._agent_event_publisher = agent_event_publisher
28+
self._propagation_credentials_repository = propagation_credentials_repository
29+
self._tcp_port_selector = tcp_port_selector
30+
self._otp_provider = otp_provider
31+
self._create_plugin = create_plugin
32+
33+
def create(self, plugin_name: str) -> SingleUsePlugin:
34+
return self._create_plugin(
35+
plugin_name=plugin_name,
36+
agent_id=self._agent_id,
37+
agent_binary_repository=self._agent_binary_repository,
38+
agent_event_publisher=self._agent_event_publisher,
39+
propagation_credentials_repository=self._propagation_credentials_repository,
40+
tcp_port_selector=self._tcp_port_selector,
41+
otp_provider=self._otp_provider,
42+
)

‎monkey/infection_monkey/plugin/i_plugin_factory.py

+19
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
from abc import ABC, abstractmethod
2+
3+
from serpentarium import SingleUsePlugin
4+
5+
6+
class IPluginFactory(ABC):
7+
"""
8+
Abstract base class for factories that create plugins for specific `AgentPluginType`s
9+
"""
10+
11+
@abstractmethod
12+
def create(self, plugin_name: str) -> SingleUsePlugin:
13+
"""
14+
Create a plugin with the given name
15+
16+
:param plugin_name: The name of the plugin to create
17+
:return: A plugin with the given name
18+
"""
19+
pass

‎monkey/infection_monkey/plugin/multiprocessing_plugin_wrapper.py

+42
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
import threading
2+
from logging import getLogger
3+
from typing import Any
4+
5+
from serpentarium import MultiUsePlugin, PluginLoader
6+
7+
logger = getLogger(__name__)
8+
9+
# NOTE: This should probably get moved to serpentarium.
10+
class MultiprocessingPluginWrapper(MultiUsePlugin):
11+
"""
12+
Wraps a MultiprocessingPlugin so it can be used like a MultiUsePlugin
13+
"""
14+
15+
process_start_lock = threading.Lock()
16+
17+
def __init__(self, *, plugin_loader: PluginLoader, plugin_name: str, **kwargs):
18+
self._plugin_loader = plugin_loader
19+
self._name = plugin_name
20+
self._constructor_kwargs = kwargs
21+
22+
def run(self, **kwargs) -> Any:
23+
logger.debug(f"Constructing a new instance of {self._name}")
24+
plugin = self._plugin_loader.load_multiprocessing_plugin(
25+
plugin_name=self._name, **self._constructor_kwargs
26+
)
27+
28+
# HERE BE DRAGONS! multiprocessing.Process.start() is not thread-safe on Linux when used
29+
# with the "spawn" method. See https://github.com/pyinstaller/pyinstaller/issues/7410 for
30+
# more details.
31+
# UPDATE: This has been resolved in PyInstaller 5.8.0. Consider removing this lock, but
32+
# leaving a comment here for future reference.
33+
with MultiprocessingPluginWrapper.process_start_lock:
34+
logger.debug("Invoking plugin.start()")
35+
plugin.start(**kwargs)
36+
37+
plugin.join()
38+
return plugin.return_value
39+
40+
@property
41+
def name(self) -> str:
42+
return self._name

‎monkey/infection_monkey/puppet/plugin_registry.py

+14-72
Original file line numberDiff line numberDiff line change
@@ -1,41 +1,28 @@
11
import logging
2-
import threading
32
from copy import copy
43
from threading import RLock
54
from typing import Any, Dict
65

7-
from serpentarium import MultiUsePlugin, PluginLoader, PluginThreadName, SingleUsePlugin
6+
from serpentarium import PluginThreadName, SingleUsePlugin
87

98
from common import OperatingSystem
109
from common.agent_plugins import AgentPlugin, AgentPluginType
11-
from common.event_queue import IAgentEventPublisher
12-
from common.types import AgentID
13-
from infection_monkey.exploit import IAgentBinaryRepository, IAgentOTPProvider
1410
from infection_monkey.i_puppet import UnknownPluginError
1511
from infection_monkey.island_api_client import IIslandAPIClient, IslandAPIRequestError
16-
from infection_monkey.network import TCPPortSelector
17-
from infection_monkey.propagation_credentials_repository import IPropagationCredentialsRepository
12+
from infection_monkey.plugin.i_plugin_factory import IPluginFactory
1813

1914
from . import PluginSourceExtractor
2015

2116
logger = logging.getLogger()
2217

2318

24-
# TODO: We should add an ExploiterPluginFactor and pass that to this component instead of passing
25-
# all of the requirements for exploiters.
2619
class PluginRegistry:
2720
def __init__(
2821
self,
2922
operating_system: OperatingSystem,
3023
island_api_client: IIslandAPIClient,
3124
plugin_source_extractor: PluginSourceExtractor,
32-
plugin_loader: PluginLoader,
33-
agent_binary_repository: IAgentBinaryRepository,
34-
agent_event_publisher: IAgentEventPublisher,
35-
propagation_credentials_repository: IPropagationCredentialsRepository,
36-
tcp_port_selector: TCPPortSelector,
37-
otp_provider: IAgentOTPProvider,
38-
agent_id: AgentID,
25+
plugin_factories: Dict[AgentPluginType, IPluginFactory],
3926
):
4027
"""
4128
`self._registry` looks like -
@@ -51,14 +38,8 @@ def __init__(
5138
self._operating_system = operating_system
5239
self._island_api_client = island_api_client
5340
self._plugin_source_extractor = plugin_source_extractor
54-
self._plugin_loader = plugin_loader
55-
self._agent_binary_repository = agent_binary_repository
56-
self._agent_event_publisher = agent_event_publisher
57-
self._propagation_credentials_repository = propagation_credentials_repository
58-
self._tcp_port_selector = tcp_port_selector
59-
self._otp_provider = otp_provider
60-
61-
self._agent_id = agent_id
41+
self._plugin_factories = plugin_factories
42+
6243
self._lock = RLock()
6344

6445
def get_plugin(self, plugin_type: AgentPluginType, plugin_name: str) -> Any:
@@ -74,18 +55,15 @@ def get_plugin(self, plugin_type: AgentPluginType, plugin_name: str) -> Any:
7455
def _load_plugin_from_island(self, plugin_name: str, plugin_type: AgentPluginType):
7556
agent_plugin = self._download_plugin_from_island(plugin_name, plugin_type)
7657
self._plugin_source_extractor.extract_plugin_source(agent_plugin)
77-
multiprocessing_plugin = MultiprocessingPluginWrapper(
78-
plugin_loader=self._plugin_loader,
79-
plugin_name=plugin_name,
80-
reset_modules_cache=False,
81-
main_thread_name=PluginThreadName.CALLING_THREAD,
82-
agent_id=self._agent_id,
83-
agent_binary_repository=self._agent_binary_repository,
84-
agent_event_publisher=self._agent_event_publisher,
85-
propagation_credentials_repository=self._propagation_credentials_repository,
86-
tcp_port_selector=self._tcp_port_selector,
87-
otp_provider=self._otp_provider,
88-
)
58+
59+
if plugin_type in self._plugin_factories:
60+
factory = self._plugin_factories[plugin_type]
61+
multiprocessing_plugin = factory.create(plugin_name)
62+
else:
63+
raise UnknownPluginError(
64+
"Loading of custom plugins has not been implemented for plugin type "
65+
f"'{plugin_type.value}'"
66+
)
8967

9068
self.load_plugin(plugin_type, plugin_name, multiprocessing_plugin)
9169

@@ -110,39 +88,3 @@ def load_plugin(self, plugin_type: AgentPluginType, plugin_name: str, plugin: ob
11088

11189
self._registry[plugin_type][plugin_name] = plugin
11290
logger.debug(f"Plugin '{plugin_name}' loaded")
113-
114-
115-
# NOTE: This should probably get moved to serpentarium.
116-
class MultiprocessingPluginWrapper(MultiUsePlugin):
117-
"""
118-
Wraps a MultiprocessingPlugin so it can be used like a MultiUsePlugin
119-
"""
120-
121-
process_start_lock = threading.Lock()
122-
123-
def __init__(self, *, plugin_loader: PluginLoader, plugin_name: str, **kwargs):
124-
self._plugin_loader = plugin_loader
125-
self._name = plugin_name
126-
self._constructor_kwargs = kwargs
127-
128-
def run(self, **kwargs) -> Any:
129-
logger.debug(f"Constructing a new instance of {self._name}")
130-
plugin = self._plugin_loader.load_multiprocessing_plugin(
131-
plugin_name=self._name, **self._constructor_kwargs
132-
)
133-
134-
# HERE BE DRAGONS! multiprocessing.Process.start() is not thread-safe on Linux when used
135-
# with the "spawn" method. See https://github.com/pyinstaller/pyinstaller/issues/7410 for
136-
# more details.
137-
# UPDATE: This has been resolved in PyInstaller 5.8.0. Consider removing this lock, but
138-
# leaving a comment here for future reference.
139-
with MultiprocessingPluginWrapper.process_start_lock:
140-
logger.debug("Invoking plugin.start()")
141-
plugin.start(**kwargs)
142-
143-
plugin.join()
144-
return plugin.return_value
145-
146-
@property
147-
def name(self) -> str:
148-
return self._name

0 commit comments

Comments
 (0)