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