forked from p15670423/monkey
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:
parent
a96b82fa0f
commit
c5c8bc1d2f
|
@ -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)
|
||||
|
|
|
@ -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)
|
Loading…
Reference in New Issue