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)
This commit is contained in:
vakarisz 2022-10-07 16:04:51 +03:00
parent a96b82fa0f
commit c5c8bc1d2f
2 changed files with 74 additions and 1 deletions

View File

@ -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)

View File

@ -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)