forked from p34709852/monkey
Refactor DocumentEncryptor class into a series of methods.
DocumentEncryptor class serves no purpose because it holds no state, sensitive_fields can be passed as a parameter to methods
This commit is contained in:
parent
f3865d022b
commit
854ce4e1e1
|
@ -3,7 +3,13 @@ from __future__ import annotations
|
|||
from bson import json_util
|
||||
from mongoengine import DictField, Document
|
||||
|
||||
from monkey_island.cc.models.utils import report_encryptor
|
||||
from monkey_island.cc.models.utils import document_encryptor
|
||||
from monkey_island.cc.models.utils.document_encryptor import SensitiveField
|
||||
from monkey_island.cc.models.utils.field_encryptors.string_list_encryptor import StringListEncryptor
|
||||
|
||||
sensitive_fields = [
|
||||
SensitiveField(path="overview.config_passwords", field_encryptor=StringListEncryptor)
|
||||
]
|
||||
|
||||
|
||||
class Report(Document):
|
||||
|
@ -18,7 +24,7 @@ class Report(Document):
|
|||
@staticmethod
|
||||
def save_report(report_dict: dict):
|
||||
report_dict = _encode_dot_char_before_mongo_insert(report_dict)
|
||||
report_dict = report_encryptor.encrypt(report_dict)
|
||||
report_dict = document_encryptor.encrypt(sensitive_fields, report_dict)
|
||||
Report.objects.delete()
|
||||
Report(
|
||||
overview=report_dict["overview"],
|
||||
|
@ -30,7 +36,9 @@ class Report(Document):
|
|||
@staticmethod
|
||||
def get_report() -> dict:
|
||||
report_dict = Report.objects.first().to_mongo()
|
||||
return _decode_dot_char_before_mongo_insert(report_encryptor.decrypt(report_dict))
|
||||
return _decode_dot_char_before_mongo_insert(
|
||||
document_encryptor.decrypt(sensitive_fields, report_dict)
|
||||
)
|
||||
|
||||
|
||||
def _encode_dot_char_before_mongo_insert(report_dict):
|
||||
|
|
|
@ -1,54 +1,48 @@
|
|||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable, List, Type
|
||||
|
||||
import dpath.util
|
||||
|
||||
from monkey_island.cc.models.utils.field_types.field_type_abc import FieldTypeABC
|
||||
from monkey_island.cc.models.utils.field_encryptors.i_field_encryptor import IFieldEncryptor
|
||||
|
||||
|
||||
class FieldNotFoundError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class SensitiveField:
|
||||
path: str
|
||||
path_separator = "."
|
||||
field_type: Type[FieldTypeABC]
|
||||
field_encryptor: Type[IFieldEncryptor]
|
||||
|
||||
|
||||
class DocumentEncryptor(ABC):
|
||||
@property
|
||||
@abstractmethod
|
||||
def sensitive_fields(self) -> List[SensitiveField]:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def encrypt(cls, document_dict: dict) -> dict:
|
||||
for sensitive_field in cls.sensitive_fields:
|
||||
DocumentEncryptor._apply_operation_to_document_field(
|
||||
document_dict, sensitive_field, sensitive_field.field_type.encrypt
|
||||
)
|
||||
|
||||
return document_dict
|
||||
|
||||
@classmethod
|
||||
def decrypt(cls, document_dict: dict) -> dict:
|
||||
for sensitive_field in cls.sensitive_fields:
|
||||
DocumentEncryptor._apply_operation_to_document_field(
|
||||
document_dict, sensitive_field, sensitive_field.field_type.decrypt
|
||||
)
|
||||
return document_dict
|
||||
|
||||
@staticmethod
|
||||
def _apply_operation_to_document_field(
|
||||
report: dict, sensitive_field: SensitiveField, operation: Callable
|
||||
):
|
||||
field_value = dpath.util.get(
|
||||
report, sensitive_field.path, sensitive_field.path_separator, None
|
||||
def encrypt(sensitive_fields: List[SensitiveField], document_dict: dict) -> dict:
|
||||
for sensitive_field in sensitive_fields:
|
||||
_apply_operation_to_document_field(
|
||||
document_dict, sensitive_field, sensitive_field.field_encryptor.encrypt
|
||||
)
|
||||
if field_value is None:
|
||||
raise Exception(
|
||||
f"Can't encrypt object because the path {sensitive_field.path} doesn't exist."
|
||||
)
|
||||
|
||||
modified_value = operation(field_value)
|
||||
return document_dict
|
||||
|
||||
dpath.util.set(report, sensitive_field.path, modified_value, sensitive_field.path_separator)
|
||||
|
||||
def decrypt(sensitive_fields: List[SensitiveField], document_dict: dict) -> dict:
|
||||
for sensitive_field in sensitive_fields:
|
||||
_apply_operation_to_document_field(
|
||||
document_dict, sensitive_field, sensitive_field.field_encryptor.decrypt
|
||||
)
|
||||
return document_dict
|
||||
|
||||
|
||||
def _apply_operation_to_document_field(
|
||||
report: dict, sensitive_field: SensitiveField, operation: Callable
|
||||
):
|
||||
field_value = dpath.util.get(report, sensitive_field.path, sensitive_field.path_separator, None)
|
||||
if field_value is None:
|
||||
raise FieldNotFoundError(
|
||||
f"Can't encrypt object because the path {sensitive_field.path} doesn't exist."
|
||||
)
|
||||
|
||||
modified_value = operation(field_value)
|
||||
|
||||
dpath.util.set(report, sensitive_field.path, modified_value, sensitive_field.path_separator)
|
||||
|
|
|
@ -4,8 +4,8 @@ from typing import List
|
|||
import pytest
|
||||
|
||||
from monkey_island.cc.models import Report
|
||||
from monkey_island.cc.models.utils.document_encryptor import SensitiveField
|
||||
from monkey_island.cc.models.utils.field_encryptors.i_field_encryptor import IFieldEncryptor
|
||||
from monkey_island.cc.models.utils.report_encryptor import SensitiveField
|
||||
|
||||
MOCK_SENSITIVE_FIELD_CONTENTS = ["the_string", "the_string2"]
|
||||
MOCK_REPORT_DICT = {
|
||||
|
@ -19,32 +19,30 @@ MOCK_REPORT_DICT = {
|
|||
}
|
||||
|
||||
|
||||
class MockFieldEncryptor(IFieldEncryptor):
|
||||
class MockStringListEncryptor(IFieldEncryptor):
|
||||
plaintext = []
|
||||
|
||||
@staticmethod
|
||||
def encrypt(value: List[str]) -> List[str]:
|
||||
return [MockFieldEncryptor._encrypt(v) for v in value]
|
||||
return [MockStringListEncryptor._encrypt(v) for v in value]
|
||||
|
||||
@staticmethod
|
||||
def _encrypt(value: str) -> str:
|
||||
MockFieldEncryptor.plaintext.append(value)
|
||||
return f"ENCRYPTED_{str(len(MockFieldEncryptor.plaintext) - 1)}"
|
||||
MockStringListEncryptor.plaintext.append(value)
|
||||
return f"ENCRYPTED_{str(len(MockStringListEncryptor.plaintext) - 1)}"
|
||||
|
||||
@staticmethod
|
||||
def decrypt(value: List[str]) -> List[str]:
|
||||
return MockFieldEncryptor.plaintext
|
||||
return MockStringListEncryptor.plaintext
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def patch_sensitive_fields(monkeypatch):
|
||||
mock_sensitive_fields = [
|
||||
SensitiveField("overview.foo.the_key", MockFieldEncryptor),
|
||||
SensitiveField("overview.bar.the_key", MockFieldEncryptor),
|
||||
SensitiveField("overview.foo.the_key", MockStringListEncryptor),
|
||||
SensitiveField("overview.bar.the_key", MockStringListEncryptor),
|
||||
]
|
||||
monkeypatch.setattr(
|
||||
"monkey_island.cc.models.utils.report_encryptor.sensitive_fields", mock_sensitive_fields
|
||||
)
|
||||
monkeypatch.setattr("monkey_island.cc.models.report.sensitive_fields", mock_sensitive_fields)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("uses_database")
|
||||
|
|
Loading…
Reference in New Issue