diff --git a/monkey/tests/unit_tests/monkey_island/cc/models/test_monkey.py b/monkey/tests/unit_tests/monkey_island/cc/models/test_monkey.py index f5a00e5e7..e25871378 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/models/test_monkey.py +++ b/monkey/tests/unit_tests/monkey_island/cc/models/test_monkey.py @@ -3,7 +3,7 @@ import uuid import pytest -from monkey_island.cc.models.monkey import Monkey, MonkeyNotFoundError +from monkey_island.cc.models.monkey import Monkey, MonkeyNotFoundError, ParentNotFoundError from monkey_island.cc.models.monkey_ttl import MonkeyTtl logger = logging.getLogger(__name__) @@ -162,3 +162,35 @@ class TestMonkey: cache_info_after_query = Monkey.is_monkey.storage.backend.cache_info() assert cache_info_after_query.hits == 2 + + @pytest.mark.usefixtures("uses_database") + def test_has_parent(self): + monkey_1 = Monkey(guid=str(uuid.uuid4())) + monkey_2 = Monkey(guid=str(uuid.uuid4())) + monkey_1.parent = [[monkey_2.guid]] + monkey_1.save() + assert monkey_1.has_parent() + + @pytest.mark.usefixtures("uses_database") + def test_has_no_parent(self): + monkey_1 = Monkey(guid=str(uuid.uuid4())) + monkey_1.parent = [[monkey_1.guid]] + monkey_1.save() + assert not monkey_1.has_parent() + + @pytest.mark.usefixtures("uses_database") + def test_get_parent(self): + monkey_1 = Monkey(guid=str(uuid.uuid4())) + monkey_2 = Monkey(guid=str(uuid.uuid4())) + monkey_1.parent = [[monkey_2.guid]] + monkey_1.save() + monkey_2.save() + assert monkey_1.get_parent().guid == monkey_2.guid + + @pytest.mark.usefixtures("uses_database") + def test_get_parent_no_parent(self): + monkey_1 = Monkey(guid=str(uuid.uuid4())) + monkey_1.parent = [[monkey_1.guid]] + monkey_1.save() + with pytest.raises(ParentNotFoundError): + monkey_1.get_parent() diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/test_infection_lifecycle.py b/monkey/tests/unit_tests/monkey_island/cc/services/test_infection_lifecycle.py new file mode 100644 index 000000000..389bf3c9c --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/services/test_infection_lifecycle.py @@ -0,0 +1,101 @@ +import uuid + +import pytest + +from monkey_island.cc.models import Config, Monkey +from monkey_island.cc.models.agent_controls import AgentControls +from monkey_island.cc.services.infection_lifecycle import should_agent_die + + +@pytest.mark.usefixtures("uses_database") +def test_should_agent_die_by_config(monkeypatch): + monkey = Monkey(guid=str(uuid.uuid4())) + monkey.config = Config(alive=False) + monkey.save() + assert should_agent_die(monkey.guid) + + monkeypatch.setattr( + "monkey_island.cc.services.infection_lifecycle._is_monkey_killed_manually", lambda _: False + ) + monkey.config.alive = True + monkey.save() + assert not should_agent_die(monkey.guid) + + +def create_monkey(launch_time): + monkey = Monkey(guid=str(uuid.uuid4())) + monkey.config = Config(alive=True) + monkey.launch_time = launch_time + monkey.save() + return monkey + + +def create_kill_event(event_time): + kill_event = AgentControls(last_stop_all=event_time) + kill_event.save() + return kill_event + + +def create_parent(child_monkey, launch_time): + monkey_parent = Monkey(guid=str(uuid.uuid4())) + child_monkey.parent = [[monkey_parent.guid]] + monkey_parent.launch_time = launch_time + monkey_parent.save() + child_monkey.save() + + +@pytest.mark.usefixtures("uses_database") +def test_was_agent_killed_manually(monkeypatch): + monkey = create_monkey(launch_time=2) + + create_kill_event(event_time=3) + + assert should_agent_die(monkey.guid) + + +@pytest.mark.usefixtures("uses_database") +def test_agent_killed_on_wakeup(monkeypatch): + monkey = create_monkey(launch_time=2) + + create_kill_event(event_time=2) + + assert should_agent_die(monkey.guid) + + +@pytest.mark.usefixtures("uses_database") +def test_manual_kill_dont_affect_new_monkeys(monkeypatch): + monkey = create_monkey(launch_time=3) + + create_kill_event(event_time=2) + + assert not should_agent_die(monkey.guid) + + +@pytest.mark.usefixtures("uses_database") +def test_parent_manually_killed(monkeypatch): + monkey = create_monkey(launch_time=3) + create_parent(child_monkey=monkey, launch_time=1) + + create_kill_event(event_time=2) + + assert should_agent_die(monkey.guid) + + +@pytest.mark.usefixtures("uses_database") +def test_parent_manually_killed_on_wakeup(monkeypatch): + monkey = create_monkey(launch_time=3) + create_parent(child_monkey=monkey, launch_time=2) + + create_kill_event(event_time=2) + + assert should_agent_die(monkey.guid) + + +@pytest.mark.usefixtures("uses_database") +def test_manual_kill_dont_affect_new_monkeys_with_parent(monkeypatch): + monkey = create_monkey(launch_time=3) + create_parent(child_monkey=monkey, launch_time=2) + + create_kill_event(event_time=1) + + assert not should_agent_die(monkey.guid)