Refactored the GenerateMapNodes function to use the DAL and now it filters windows nodes

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'],
+                'id': monkey.guid,
-                'label': '{0} : {1}'.format(monkey['hostname'], monkey['ip_addresses'][0]),
+                'label': '{0} : {1}'.format(monkey.hostname, monkey.ip_addresses[0]),
-                'group': 'critical' if monkey.get('critical_services', []) else 'normal',
+                'group': 'critical' if monkey.critical_services is not None else 'normal',
-                'services': monkey.get('critical_services', []),
+                'services': monkey.critical_services,
-                'hostname': monkey['hostname']
+                '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)
+
+
+