From 712ce4622d145c47fc3c38e1527f03dc1ab6fa6e Mon Sep 17 00:00:00 2001 From: Shay Nehmad Date: Wed, 5 Jun 2019 18:28:54 +0300 Subject: [PATCH] Refactored the GenerateMapNodes function to use the DAL and now it filters windows nodes --- .../monkey_island/cc/services/pth_report.py | 13 ++-- .../cc/services/test_PTHReportService.py | 69 +++++++++++++++++++ 2 files changed, 76 insertions(+), 6 deletions(-) create mode 100644 monkey/monkey_island/cc/services/test_PTHReportService.py diff --git a/monkey/monkey_island/cc/services/pth_report.py b/monkey/monkey_island/cc/services/pth_report.py index 93fd51989..9f3b9769f 100644 --- a/monkey/monkey_island/cc/services/pth_report.py +++ b/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 ] diff --git a/monkey/monkey_island/cc/services/test_PTHReportService.py b/monkey/monkey_island/cc/services/test_PTHReportService.py new file mode 100644 index 000000000..24e560ff0 --- /dev/null +++ b/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) + + +