From ccadfccf5e20ca4248fdeba66b3a16a610ef6136 Mon Sep 17 00:00:00 2001
From: Kekoa Kaaikala <kekoa.kaaikala@gmail.com>
Date: Mon, 19 Sep 2022 14:45:49 +0000
Subject: [PATCH] Island: Remove duplicate encryption code

---
 .../cc/repository/agent_event_encryption.py   | 14 ++---
 .../mongo_agent_event_repository.py           | 51 ++-----------------
 2 files changed, 11 insertions(+), 54 deletions(-)

diff --git a/monkey/monkey_island/cc/repository/agent_event_encryption.py b/monkey/monkey_island/cc/repository/agent_event_encryption.py
index dd0d813cf..8904773ed 100644
--- a/monkey/monkey_island/cc/repository/agent_event_encryption.py
+++ b/monkey/monkey_island/cc/repository/agent_event_encryption.py
@@ -34,13 +34,14 @@ def encrypt_event(
     if not isinstance(event_data, dict):
         raise TypeError("Event encryption only supported for dict")
 
+    data = event_data.copy()
     for field in fields:
-        event_data[ENCRYPTED_PREFIX + field] = str(
+        data[ENCRYPTED_PREFIX + field] = str(
             encrypt(json.dumps(event_data[field]).encode()), "utf-8"
         )
-        del event_data[field]
+        del data[field]
 
-    return event_data
+    return data
 
 
 def decrypt_event(
@@ -57,11 +58,12 @@ def decrypt_event(
     if not isinstance(event_data, dict):
         raise TypeError("Event decryption only supported for dict")
 
+    data = event_data.copy()
     for field in event_data.keys():
         if field.startswith("encrypted_"):
-            event_data[field[len(ENCRYPTED_PREFIX) :]] = json.loads(
+            data[field[len(ENCRYPTED_PREFIX) :]] = json.loads(
                 str(decrypt(event_data[field].encode()), "utf-8")
             )
-            del event_data[field]
+            del data[field]
 
-    return event_data
+    return data
diff --git a/monkey/monkey_island/cc/repository/mongo_agent_event_repository.py b/monkey/monkey_island/cc/repository/mongo_agent_event_repository.py
index 336039319..0194f3d3d 100644
--- a/monkey/monkey_island/cc/repository/mongo_agent_event_repository.py
+++ b/monkey/monkey_island/cc/repository/mongo_agent_event_repository.py
@@ -1,62 +1,17 @@
-import json
-from typing import Any, Callable, Dict, Iterable, MutableMapping, Sequence, Type
+from typing import Any, Dict, MutableMapping, Sequence, Type
 
 from pymongo import MongoClient
 
-from common.agent_event_serializers import (
-    EVENT_TYPE_FIELD,
-    AgentEventSerializerRegistry,
-    JSONSerializable,
-)
+from common.agent_event_serializers import EVENT_TYPE_FIELD, AgentEventSerializerRegistry
 from common.agent_events import AbstractAgentEvent
 from common.types import AgentID
 from monkey_island.cc.repository import IAgentEventRepository
 from monkey_island.cc.server_utils.encryption import ILockableEncryptor
 
 from . import RemovalError, RetrievalError, StorageError
+from .agent_event_encryption import decrypt_event, encrypt_event, get_fields_to_encrypt
 from .consts import MONGO_OBJECT_ID_KEY
 
-ENCRYPTED_PREFIX = "encrypted_"
-
-
-def get_fields_to_encrypt(event: AbstractAgentEvent):
-    return set(vars(AbstractAgentEvent)["__fields__"].keys()) ^ set(event.dict().keys())
-
-
-def encrypt_event(
-    encrypt: Callable[[bytes], bytes],
-    event_data: JSONSerializable,
-    fields: Iterable[str] = [],
-) -> JSONSerializable:
-    if not isinstance(event_data, dict):
-        raise TypeError("Event encryption only supported for dict")
-
-    data = event_data.copy()
-    for field in fields:
-        data[ENCRYPTED_PREFIX + field] = str(
-            encrypt(json.dumps(event_data[field]).encode()), "utf-8"
-        )
-        del data[field]
-
-    return data
-
-
-def decrypt_event(
-    decrypt: Callable[[bytes], bytes], event_data: JSONSerializable
-) -> JSONSerializable:
-    if not isinstance(event_data, dict):
-        raise TypeError("Event decryption only supported for dict")
-
-    data = event_data.copy()
-    for field in event_data.keys():
-        if field.startswith("encrypted_"):
-            data[field[len(ENCRYPTED_PREFIX) :]] = json.loads(
-                str(decrypt(event_data[field].encode()), "utf-8")
-            )
-            del data[field]
-
-    return data
-
 
 class MongoAgentEventRepository(IAgentEventRepository):
     """A repository for storing and retrieving events in MongoDB"""