Merge pull request #338 from guardicore/improvement/232-dont-show-linux-in-cred-map

[CR] Improvement/232 dont show linux in cred map

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

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())))

monkey/monkey_island/cc/services/pth_report.py

@@ -1,6 +1,7 @@
 from itertools import product
 
 from monkey_island.cc.database import mongo
+from monkey_island.cc.models import Monkey
 from bson import ObjectId
 
 from monkey_island.cc.services.groups_and_users_consts import USERTYPE
@@ -216,15 +217,15 @@ class PTHReportService(object):
 
     @staticmethod
     def generate_map_nodes():
-        monkeys = mongo.db.monkey.find({}, {'_id': 1, 'hostname': 1, 'critical_services': 1, 'ip_addresses': 1})
+        monkeys = filter(lambda m: m.get_os() == "windows", Monkey.objects())
 
         return [
             {
-                'id': monkey['_id'],
-                'label': '{0} : {1}'.format(monkey['hostname'], monkey['ip_addresses'][0]),
-                'group': 'critical' if monkey.get('critical_services', []) else 'normal',
-                'services': monkey.get('critical_services', []),
-                'hostname': monkey['hostname']
+                'id': monkey.guid,
+                'label': '{0} : {1}'.format(monkey.hostname, monkey.ip_addresses[0]),
+                'group': 'critical' if monkey.critical_services is not None else 'normal',
+                'services': monkey.critical_services,
+                'hostname': monkey.hostname
             } for monkey in monkeys
         ]
 

monkey/monkey_island/cc/services/test_PTHReportService.py

@@ -0,0 +1,69 @@
+import uuid
+
+from monkey_island.cc.models import Monkey
+from monkey_island.cc.services.pth_report import PTHReportService
+from monkey_island.cc.testing.IslandTestCase import IslandTestCase
+
+
+class TestPTHReportServiceGenerateMapNodes(IslandTestCase):
+    def test_generate_map_nodes(self):
+        self.fail_if_not_testing_env()
+        self.clean_monkey_db()
+
+        self.assertEqual(PTHReportService.generate_map_nodes(), [])
+
+        windows_monkey_with_services = Monkey(
+            guid=str(uuid.uuid4()),
+            hostname="A_Windows_PC_1",
+            critical_services=["aCriticalService", "Domain Controller"],
+            ip_addresses=["1.1.1.1", "2.2.2.2"],
+            description="windows 10"
+        )
+        windows_monkey_with_services.save()
+
+        windows_monkey_with_no_services = Monkey(
+            guid=str(uuid.uuid4()),
+            hostname="A_Windows_PC_2",
+            critical_services=[],
+            ip_addresses=["3.3.3.3"],
+            description="windows 10"
+        )
+        windows_monkey_with_no_services.save()
+
+        linux_monkey = Monkey(
+            guid=str(uuid.uuid4()),
+            hostname="A_Linux_PC",
+            ip_addresses=["4.4.4.4"],
+            description="linux ubuntu"
+        )
+        linux_monkey.save()
+
+        map_nodes = PTHReportService.generate_map_nodes()
+
+        self.assertEquals(2, len(map_nodes))
+
+    def test_generate_map_nodes_parsing(self):
+        self.fail_if_not_testing_env()
+        self.clean_monkey_db()
+
+        monkey_id = str(uuid.uuid4())
+        hostname = "A_Windows_PC_1"
+        windows_monkey_with_services = Monkey(
+            guid=monkey_id,
+            hostname=hostname,
+            critical_services=["aCriticalService", "Domain Controller"],
+            ip_addresses=["1.1.1.1"],
+            description="windows 10"
+        )
+        windows_monkey_with_services.save()
+
+        map_nodes = PTHReportService.generate_map_nodes()
+
+        self.assertEquals(map_nodes[0]["id"], monkey_id)
+        self.assertEquals(map_nodes[0]["label"], "A_Windows_PC_1 : 1.1.1.1")
+        self.assertEquals(map_nodes[0]["group"], "critical")
+        self.assertEquals(len(map_nodes[0]["services"]), 2)
+        self.assertEquals(map_nodes[0]["hostname"], hostname)
+
+
+

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()