From fa8736e2fed40dda766e09954b3d82d2be8ec227 Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
Date: Wed, 14 Sep 2022 09:05:42 -0400
Subject: [PATCH] Island: Omit the mongo object ID from Machine query results

---
 .../cc/repository/mongo_machine_repository.py   | 17 +++++++----------
 vulture_allowlist.py                            |  2 ++
 2 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/monkey/monkey_island/cc/repository/mongo_machine_repository.py b/monkey/monkey_island/cc/repository/mongo_machine_repository.py
index 4d1a36470..fab038694 100644
--- a/monkey/monkey_island/cc/repository/mongo_machine_repository.py
+++ b/monkey/monkey_island/cc/repository/mongo_machine_repository.py
@@ -1,6 +1,6 @@
 from ipaddress import IPv4Address
 from threading import Lock
-from typing import Any, MutableMapping, Sequence
+from typing import Any, Sequence
 
 from pymongo import MongoClient
 
@@ -58,36 +58,33 @@ class MongoMachineRepository(IMachineRepository):
 
     def _find_one(self, key: str, search_value: Any) -> Machine:
         try:
-            machine_dict = self._machines_collection.find_one({key: search_value})
+            machine_dict = self._machines_collection.find_one(
+                {key: search_value}, {MONGO_OBJECT_ID_KEY: False}
+            )
         except Exception as err:
             raise RetrievalError(f'Error retrieving machine with "{key} == {search_value}": {err}')
 
         if machine_dict is None:
             raise UnknownRecordError(f'Unknown machine with "{key} == {search_value}"')
 
-        return MongoMachineRepository._mongo_record_to_machine(machine_dict)
+        return Machine(**machine_dict)
 
     def get_machines_by_ip(self, ip: IPv4Address) -> Sequence[Machine]:
         ip_regex = "^" + str(ip).replace(".", "\\.") + "\\/.*$"
         query = {"network_interfaces": {"$elemMatch": {"$regex": ip_regex}}}
 
         try:
-            cursor = self._machines_collection.find(query)
+            cursor = self._machines_collection.find(query, {MONGO_OBJECT_ID_KEY: False})
         except Exception as err:
             raise RetrievalError(f'Error retrieving machines with ip "{ip}": {err}')
 
-        machines = list(map(MongoMachineRepository._mongo_record_to_machine, cursor))
+        machines = list(map(lambda m: Machine(**m), cursor))
 
         if len(machines) == 0:
             raise UnknownRecordError(f'No machines found with IP "{ip}"')
 
         return machines
 
-    @staticmethod
-    def _mongo_record_to_machine(mongo_record: MutableMapping[str, Any]) -> Machine:
-        del mongo_record[MONGO_OBJECT_ID_KEY]
-        return Machine(**mongo_record)
-
     def reset(self):
         try:
             self._machines_collection.drop()
diff --git a/vulture_allowlist.py b/vulture_allowlist.py
index c77390784..c4c33b662 100644
--- a/vulture_allowlist.py
+++ b/vulture_allowlist.py
@@ -12,6 +12,7 @@ from infection_monkey.exploit.log4shell_utils.ldap_server import LDAPServerFacto
 from monkey_island.cc.event_queue import IslandEventTopic, PyPubSubIslandEventQueue
 from monkey_island.cc.models import Report
 from monkey_island.cc.models.networkmap import Arc, NetworkMap
+from monkey_island.cc.repository import MongoMachineRepository
 from monkey_island.cc.repository.attack.IMitigationsRepository import IMitigationsRepository
 from monkey_island.cc.repository.i_agent_repository import IAgentRepository
 from monkey_island.cc.repository.i_attack_repository import IAttackRepository
@@ -277,6 +278,7 @@ ICredentialsRepository.save_configured_credentials
 ITelemetryRepository.get_telemetries
 IEventRepository.get_events
 IFindingRepository.get_findings
+MongoMachineRepository
 key_list
 simulation
 netmap