diff --git a/monkey/monkey_island/cc/models/monkey.py b/monkey/monkey_island/cc/models/monkey.py index bb018caa9..3e1e3d7c5 100644 --- a/monkey/monkey_island/cc/models/monkey.py +++ b/monkey/monkey_island/cc/models/monkey.py @@ -55,6 +55,14 @@ class Monkey(Document): monkey_is_dead = True return monkey_is_dead + def get_os(self): + os = "unknown" + if self.description.lower().find("linux") != -1: + os = "linux" + elif self.description.lower().find("windows") != -1: + os = "windows" + return os + class MonkeyNotFoundError(Exception): pass diff --git a/monkey/monkey_island/cc/models/test_monkey.py b/monkey/monkey_island/cc/models/test_monkey.py index 008fb0ce6..a744db6b6 100644 --- a/monkey/monkey_island/cc/models/test_monkey.py +++ b/monkey/monkey_island/cc/models/test_monkey.py @@ -1,13 +1,13 @@ import uuid from time import sleep -from unittest import TestCase from monkey import Monkey from monkey_island.cc.models.monkey import MonkeyNotFoundError +from monkey_island.cc.testing.IslandTestCase import IslandTestCase from monkey_ttl import MonkeyTtl -class TestMonkey(TestCase): +class TestMonkey(IslandTestCase): """ Make sure to set server environment to `testing` in server.json! Otherwise this will mess up your mongo instance and won't work. @@ -15,7 +15,11 @@ class TestMonkey(TestCase): Also, the working directory needs to be the working directory from which you usually run the island so the server.json file is found and loaded. """ + def test_is_dead(self): + self.fail_if_not_testing_env() + self.clean_monkey_db() + # Arrange alive_monkey_ttl = MonkeyTtl.create_ttl_expire_in(30) alive_monkey_ttl.save() @@ -43,6 +47,9 @@ class TestMonkey(TestCase): self.assertFalse(alive_monkey.is_dead()) def test_get_single_monkey_by_id(self): + self.fail_if_not_testing_env() + self.clean_monkey_db() + # Arrange a_monkey = Monkey(guid=str(uuid.uuid4())) a_monkey.save() @@ -52,3 +59,21 @@ class TestMonkey(TestCase): self.assertIsNotNone(Monkey.get_single_monkey_by_id(a_monkey.id)) # Raise on non-existent monkey self.assertRaises(MonkeyNotFoundError, Monkey.get_single_monkey_by_id, "abcdefabcdefabcdefabcdef") + + def test_get_os(self): + self.fail_if_not_testing_env() + self.clean_monkey_db() + + linux_monkey = Monkey(guid=str(uuid.uuid4()), + description="Linux shay-Virtual-Machine 4.15.0-50-generic #54-Ubuntu SMP Mon May 6 18:46:08 UTC 2019 x86_64 x86_64") + windows_monkey = Monkey(guid=str(uuid.uuid4()), + description="Windows bla bla bla") + unknown_monkey = Monkey(guid=str(uuid.uuid4()), + description="bla bla bla") + linux_monkey.save() + windows_monkey.save() + unknown_monkey.save() + + self.assertEquals(1, len(filter(lambda m: m.get_os() == "windows", Monkey.objects()))) + self.assertEquals(1, len(filter(lambda m: m.get_os() == "linux", Monkey.objects()))) + self.assertEquals(1, len(filter(lambda m: m.get_os() == "unknown", Monkey.objects()))) diff --git a/monkey/monkey_island/cc/testing/IslandTestCase.py b/monkey/monkey_island/cc/testing/IslandTestCase.py new file mode 100644 index 000000000..e894f13df --- /dev/null +++ b/monkey/monkey_island/cc/testing/IslandTestCase.py @@ -0,0 +1,12 @@ +import unittest +from monkey_island.cc.environment.environment import env +from monkey_island.cc.models import Monkey + + +class IslandTestCase(unittest.TestCase): + def fail_if_not_testing_env(self): + self.failIf(not env.testing, "Change server_config.json to testing environment.") + + @staticmethod + def clean_monkey_db(): + Monkey.objects().delete() diff --git a/monkey/monkey_island/cc/testing/__init__.py b/monkey/monkey_island/cc/testing/__init__.py new file mode 100644 index 000000000..e69de29bb