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 import platform
from socket import gethostname from socket import gethostname
from typing import Any, Mapping
from uuid import getnode from uuid import getnode
from common import OperatingSystem from common import OperatingSystem
from common.network.network_utils import get_network_interfaces from common.network.network_utils import get_network_interfaces
from monkey_island.cc.models import Machine from monkey_island.cc.models import Machine
from . import IMachineRepository, UnknownRecordError from . import IMachineRepository, StorageError, UnknownRecordError
def initialize_machine_repository(machine_repository: IMachineRepository): def initialize_machine_repository(machine_repository: IMachineRepository):
@ -33,3 +35,34 @@ def initialize_machine_repository(machine_repository: IMachineRepository):
hostname=gethostname(), hostname=gethostname(),
) )
machine_repository.upsert_machine(machine) 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)