p59342861
/
monkey
forked from p15670423/monkey
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Island: Remove duplicate encryption code

This commit is contained in:
Kekoa Kaaikala 2022-09-19 14:45:49 +00:00
parent f635c2cd5f
commit ccadfccf5e
2 changed files with 11 additions and 54 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
monkey/monkey_island/cc/repository/agent_event_encryption.py
View File

@ -34,13 +34,14 @@ def encrypt_event(
if not isinstance(event_data, dict):
raise TypeError("Event encryption only supported for dict")
data = event_data.copy()
for field in fields:
event_data[ENCRYPTED_PREFIX + field] = str(
data[ENCRYPTED_PREFIX + field] = str(
encrypt(json.dumps(event_data[field]).encode()), "utf-8"
)
del event_data[field]
del data[field]
return event_data
return data
def decrypt_event(
@ -57,11 +58,12 @@ def decrypt_event(
if not isinstance(event_data, dict):
raise TypeError("Event decryption only supported for dict")
data = event_data.copy()
for field in event_data.keys():
if field.startswith("encrypted_"):
event_data[field[len(ENCRYPTED_PREFIX) :]] = json.loads(
data[field[len(ENCRYPTED_PREFIX) :]] = json.loads(
str(decrypt(event_data[field].encode()), "utf-8")
)
del event_data[field]
del data[field]
return event_data
return data

51
monkey/monkey_island/cc/repository/mongo_agent_event_repository.py
View File

@ -1,62 +1,17 @@
import json
from typing import Any, Callable, Dict, Iterable, MutableMapping, Sequence, Type
from typing import Any, Dict, MutableMapping, Sequence, Type
from pymongo import MongoClient
from common.agent_event_serializers import (
EVENT_TYPE_FIELD,
AgentEventSerializerRegistry,
JSONSerializable,
)
from common.agent_event_serializers import EVENT_TYPE_FIELD, AgentEventSerializerRegistry
from common.agent_events import AbstractAgentEvent
from common.types import AgentID
from monkey_island.cc.repository import IAgentEventRepository
from monkey_island.cc.server_utils.encryption import ILockableEncryptor
from . import RemovalError, RetrievalError, StorageError
from .agent_event_encryption import decrypt_event, encrypt_event, get_fields_to_encrypt
from .consts import MONGO_OBJECT_ID_KEY
ENCRYPTED_PREFIX = "encrypted_"
def get_fields_to_encrypt(event: AbstractAgentEvent):
return set(vars(AbstractAgentEvent)["__fields__"].keys()) ^ set(event.dict().keys())
def encrypt_event(
encrypt: Callable[[bytes], bytes],
event_data: JSONSerializable,
fields: Iterable[str] = [],
) -> JSONSerializable:
if not isinstance(event_data, dict):
raise TypeError("Event encryption only supported for dict")
data = event_data.copy()
for field in fields:
data[ENCRYPTED_PREFIX + field] = str(
encrypt(json.dumps(event_data[field]).encode()), "utf-8"
)
del data[field]
return data
def decrypt_event(
decrypt: Callable[[bytes], bytes], event_data: JSONSerializable
) -> JSONSerializable:
if not isinstance(event_data, dict):
raise TypeError("Event decryption only supported for dict")
data = event_data.copy()
for field in event_data.keys():
if field.startswith("encrypted_"):
data[field[len(ENCRYPTED_PREFIX) :]] = json.loads(
str(decrypt(event_data[field].encode()), "utf-8")
)
del data[field]
return data
class MongoAgentEventRepository(IAgentEventRepository):
"""A repository for storing and retrieving events in MongoDB"""