forked from p15670423/monkey
UT: Use IAgentEventQueue in place of IEventQueue
This commit is contained in:
parent
acb049f1f1
commit
28cf860aeb
|
@ -6,7 +6,7 @@ from uuid import UUID
|
|||
import pytest
|
||||
from pubsub.core import Publisher
|
||||
|
||||
from common.event_queue import EventSubscriber, IEventQueue, PyPubSubEventQueue
|
||||
from common.event_queue import EventSubscriber, IAgentEventQueue, PyPubSubEventQueue
|
||||
from common.events import AbstractAgentEvent
|
||||
|
||||
EVENT_TAG_1 = "event tag 1"
|
||||
|
@ -32,7 +32,7 @@ class TestEvent2(AbstractAgentEvent):
|
|||
|
||||
|
||||
@pytest.fixture
|
||||
def event_queue() -> IEventQueue:
|
||||
def event_queue() -> IAgentEventQueue:
|
||||
return PyPubSubEventQueue(Publisher())
|
||||
|
||||
|
||||
|
@ -50,7 +50,7 @@ def event_queue_subscriber() -> Callable[[AbstractAgentEvent], None]:
|
|||
return fn
|
||||
|
||||
|
||||
def test_subscribe_all(event_queue: IEventQueue, event_queue_subscriber: EventSubscriber):
|
||||
def test_subscribe_all(event_queue: IAgentEventQueue, event_queue_subscriber: EventSubscriber):
|
||||
event_queue.subscribe_all_events(event_queue_subscriber)
|
||||
|
||||
event_queue.publish(TestEvent1(tags=frozenset({EVENT_TAG_1, EVENT_TAG_2})))
|
||||
|
@ -65,7 +65,7 @@ def test_subscribe_all(event_queue: IEventQueue, event_queue_subscriber: EventSu
|
|||
|
||||
@pytest.mark.parametrize("type_to_subscribe", [TestEvent1, TestEvent2])
|
||||
def test_subscribe_types(
|
||||
event_queue: IEventQueue, event_queue_subscriber: EventSubscriber, type_to_subscribe
|
||||
event_queue: IAgentEventQueue, event_queue_subscriber: EventSubscriber, type_to_subscribe
|
||||
):
|
||||
event_queue.subscribe_type(type_to_subscribe, event_queue_subscriber)
|
||||
|
||||
|
@ -77,7 +77,7 @@ def test_subscribe_types(
|
|||
|
||||
|
||||
def test_subscribe_tags_single_type(
|
||||
event_queue: IEventQueue, event_queue_subscriber: EventSubscriber
|
||||
event_queue: IAgentEventQueue, event_queue_subscriber: EventSubscriber
|
||||
):
|
||||
event_queue.subscribe_tag(EVENT_TAG_1, event_queue_subscriber)
|
||||
|
||||
|
@ -91,7 +91,7 @@ def test_subscribe_tags_single_type(
|
|||
|
||||
|
||||
def test_subscribe_tags_multiple_types(
|
||||
event_queue: IEventQueue, event_queue_subscriber: EventSubscriber
|
||||
event_queue: IAgentEventQueue, event_queue_subscriber: EventSubscriber
|
||||
):
|
||||
event_queue.subscribe_tag(EVENT_TAG_2, event_queue_subscriber)
|
||||
|
||||
|
@ -105,7 +105,7 @@ def test_subscribe_tags_multiple_types(
|
|||
assert {EVENT_TAG_1, EVENT_TAG_2}.issubset(event_queue_subscriber.call_tags)
|
||||
|
||||
|
||||
def test_type_tag_collision(event_queue: IEventQueue, event_queue_subscriber: EventSubscriber):
|
||||
def test_type_tag_collision(event_queue: IAgentEventQueue, event_queue_subscriber: EventSubscriber):
|
||||
event_queue.subscribe_type(TestEvent1, event_queue_subscriber)
|
||||
|
||||
event_queue.publish(TestEvent2(tags=frozenset({TestEvent1.__name__})))
|
||||
|
@ -113,7 +113,7 @@ def test_type_tag_collision(event_queue: IEventQueue, event_queue_subscriber: Ev
|
|||
assert event_queue_subscriber.call_count == 0
|
||||
|
||||
|
||||
def test_keep_subscriber_in_scope(event_queue: IEventQueue):
|
||||
def test_keep_subscriber_in_scope(event_queue: IAgentEventQueue):
|
||||
class MyCallable:
|
||||
called = False
|
||||
|
||||
|
|
|
@ -4,7 +4,7 @@ from unittest.mock import MagicMock
|
|||
import pytest
|
||||
|
||||
from common.credentials import Credentials, LMHash, NTHash, Password, Username
|
||||
from common.event_queue import IEventQueue
|
||||
from common.event_queue import IAgentEventQueue
|
||||
from common.events import CredentialsStolenEvent
|
||||
from infection_monkey.credential_collectors import MimikatzCredentialCollector
|
||||
from infection_monkey.credential_collectors.mimikatz_collector.mimikatz_credential_collector import ( # noqa: E501
|
||||
|
@ -24,7 +24,7 @@ def patch_pypykatz(win_creds: [WindowsCredentials], monkeypatch):
|
|||
|
||||
|
||||
def collect_credentials() -> Sequence[Credentials]:
|
||||
mock_event_queue = MagicMock(spec=IEventQueue)
|
||||
mock_event_queue = MagicMock(spec=IAgentEventQueue)
|
||||
return MimikatzCredentialCollector(mock_event_queue).collect_credentials()
|
||||
|
||||
|
||||
|
@ -120,7 +120,7 @@ def test_pypykatz_result_parsing_no_secrets(monkeypatch):
|
|||
|
||||
|
||||
def test_mimikatz_credentials_stolen_event_published(monkeypatch):
|
||||
mock_event_queue = MagicMock(spec=IEventQueue)
|
||||
mock_event_queue = MagicMock(spec=IAgentEventQueue)
|
||||
patch_pypykatz([], monkeypatch)
|
||||
|
||||
mimikatz_credential_collector = MimikatzCredentialCollector(mock_event_queue)
|
||||
|
@ -134,7 +134,7 @@ def test_mimikatz_credentials_stolen_event_published(monkeypatch):
|
|||
|
||||
|
||||
def test_mimikatz_credentials_stolen_event_tags(monkeypatch):
|
||||
mock_event_queue = MagicMock(spec=IEventQueue)
|
||||
mock_event_queue = MagicMock(spec=IAgentEventQueue)
|
||||
patch_pypykatz([], monkeypatch)
|
||||
|
||||
mimikatz_credential_collector = MimikatzCredentialCollector(mock_event_queue)
|
||||
|
@ -146,7 +146,7 @@ def test_mimikatz_credentials_stolen_event_tags(monkeypatch):
|
|||
|
||||
|
||||
def test_mimikatz_credentials_stolen_event_stolen_credentials(monkeypatch):
|
||||
mock_event_queue = MagicMock(spec=IEventQueue)
|
||||
mock_event_queue = MagicMock(spec=IAgentEventQueue)
|
||||
win_creds = [
|
||||
WindowsCredentials(
|
||||
username="user2", password="secret2", lm_hash="0182BD0BD4444BF8FC83B5D9042EED2E"
|
||||
|
|
|
@ -3,7 +3,7 @@ from unittest.mock import MagicMock
|
|||
import pytest
|
||||
|
||||
from common.credentials import Credentials, SSHKeypair, Username
|
||||
from common.event_queue import IEventQueue
|
||||
from common.event_queue import IAgentEventQueue
|
||||
from infection_monkey.credential_collectors import SSHCredentialCollector
|
||||
|
||||
|
||||
|
@ -25,7 +25,7 @@ def patch_ssh_handler(ssh_creds, monkeypatch):
|
|||
def test_ssh_credentials_empty_results(monkeypatch, ssh_creds, patch_telemetry_messenger):
|
||||
patch_ssh_handler(ssh_creds, monkeypatch)
|
||||
collected = SSHCredentialCollector(
|
||||
patch_telemetry_messenger, MagicMock(spec=IEventQueue)
|
||||
patch_telemetry_messenger, MagicMock(spec=IAgentEventQueue)
|
||||
).collect_credentials()
|
||||
assert not collected
|
||||
|
||||
|
@ -71,6 +71,6 @@ def test_ssh_info_result_parsing(monkeypatch, patch_telemetry_messenger):
|
|||
Credentials(identity=None, secret=ssh_keypair3),
|
||||
]
|
||||
collected = SSHCredentialCollector(
|
||||
patch_telemetry_messenger, MagicMock(spec=IEventQueue)
|
||||
patch_telemetry_messenger, MagicMock(spec=IAgentEventQueue)
|
||||
).collect_credentials()
|
||||
assert expected == collected
|
||||
|
|
Loading…
Reference in New Issue