From c5c8bc1d2fd830d5a00610032e15e908a8c23911 Mon Sep 17 00:00:00 2001
From: vakarisz <vakarisz@yahoo.com>
Date: Fri, 7 Oct 2022 16:04:51 +0300
Subject: [PATCH] Island: Add mongo_dot_encoder to encode "." characters

This encoder will be needed in mongo repository, because mongodb can't handle keys with "." character (until version 5)
---
 monkey/monkey_island/cc/repository/utils.py   | 35 +++++++++++++++-
 .../monkey_island/cc/models/test_utils.py     | 40 +++++++++++++++++++
 2 files changed, 74 insertions(+), 1 deletion(-)
 create mode 100644 monkey/tests/unit_tests/monkey_island/cc/models/test_utils.py

diff --git a/monkey/monkey_island/cc/repository/utils.py b/monkey/monkey_island/cc/repository/utils.py
index 7f11c0b46..effefe144 100644
--- a/monkey/monkey_island/cc/repository/utils.py
+++ b/monkey/monkey_island/cc/repository/utils.py
@@ -1,12 +1,14 @@
+import json
 import platform
 from socket import gethostname
+from typing import Any, Mapping
 from uuid import getnode
 
 from common import OperatingSystem
 from common.network.network_utils import get_network_interfaces
 from monkey_island.cc.models import Machine
 
-from . import IMachineRepository, UnknownRecordError
+from . import IMachineRepository, StorageError, UnknownRecordError
 
 
 def initialize_machine_repository(machine_repository: IMachineRepository):
@@ -33,3 +35,34 @@ def initialize_machine_repository(machine_repository: IMachineRepository):
             hostname=gethostname(),
         )
         machine_repository.upsert_machine(machine)
+
+
+DOT_REPLACEMENT = ",,,"
+
+
+def mongo_dot_encoder(mapping: Mapping[str, Any]) -> Mapping[str, Any]:
+    """
+    Mongo can't store keys with "." symbols (like IP's and filenames). This method
+    replaces all occurances of "." with ",,,"
+    :param mapping: Mapping to be converted to mongo compatible mapping
+    :return: Mongo compatible mapping
+    """
+    mapping_json = json.dumps(mapping)
+    if DOT_REPLACEMENT in mapping_json:
+        raise StorageError(
+            f"Mapping {mapping} already contains {DOT_REPLACEMENT}."
+            f" Aborting the encoding procedure"
+        )
+    encoded_json = mapping_json.replace(".", DOT_REPLACEMENT)
+    return json.loads(encoded_json)
+
+
+def mongo_dot_decoder(mapping: Mapping[str, Any]):
+    """
+    Mongo can't store keys with "." symbols (like IP's and filenames). This method
+    reverts changes made by "mongo_dot_encoder" by replacing all occurances of ",,," with "."
+    :param mapping: Mapping to be converted from mongo compatible mapping to original mapping
+    :return: Original mapping
+    """
+    report_as_json = json.dumps(mapping).replace(DOT_REPLACEMENT, ".")
+    return json.loads(report_as_json)
diff --git a/monkey/tests/unit_tests/monkey_island/cc/models/test_utils.py b/monkey/tests/unit_tests/monkey_island/cc/models/test_utils.py
new file mode 100644
index 000000000..33133058c
--- /dev/null
+++ b/monkey/tests/unit_tests/monkey_island/cc/models/test_utils.py
@@ -0,0 +1,40 @@
+import pytest
+
+from monkey_island.cc.repository import StorageError
+from monkey_island.cc.repository.utils import DOT_REPLACEMENT, mongo_dot_decoder, mongo_dot_encoder
+
+DATASET = [
+    ({"no:changes;expectes": "Nothing'$ changed"}, {"no:changes;expectes": "Nothing'$ changed"}),
+    (
+        {"192.168.56.1": "monkeys-running-wild.com"},
+        {
+            f"192{DOT_REPLACEMENT}168{DOT_REPLACEMENT}56{DOT_REPLACEMENT}1": f"monkeys-running-wild{DOT_REPLACEMENT}com"
+        },
+    ),
+    (
+        {"...dots...": ",comma,comma,,comedy"},
+        {
+            f"{DOT_REPLACEMENT}{DOT_REPLACEMENT}{DOT_REPLACEMENT}dots"
+            f"{DOT_REPLACEMENT}{DOT_REPLACEMENT}{DOT_REPLACEMENT}": ",comma,comma,,comedy"
+        },
+    ),
+    (
+        {"one": {"two": {"three": "this.is.nested"}}},
+        {"one": {"two": {"three": f"this{DOT_REPLACEMENT}is{DOT_REPLACEMENT}nested"}}},
+    ),
+]
+
+# This dict already contains the replacement used, encoding procedure would lose data
+FLAWED_DICT = {"one": {".two": {"three": f"this is with {DOT_REPLACEMENT} already!!!!"}}}
+
+
+@pytest.mark.parametrize("input, expected_output", DATASET)
+def test_mongo_dot_encoding_and_decoding(input, expected_output):
+    encoded = mongo_dot_encoder(input)
+    assert encoded == expected_output
+    assert mongo_dot_decoder(encoded) == input
+
+
+def test_mongo_dot_encoding__data_loss():
+    with pytest.raises(StorageError):
+        mongo_dot_encoder(FLAWED_DICT)