forked from p15670423/monkey
Refactored json parsing out of encryption/decryption functionality.
This commit is contained in:
parent
3450b80a82
commit
5c7bab7a0d
|
@ -5,7 +5,7 @@ from flask import request
|
|||
|
||||
from monkey_island.cc.resources.auth.auth import jwt_required
|
||||
from monkey_island.cc.services.config import ConfigService
|
||||
from monkey_island.cc.services.utils.config_encryption import encrypt_config
|
||||
from monkey_island.cc.services.utils.encryption import encrypt_string
|
||||
|
||||
|
||||
class ConfigurationExport(flask_restful.Resource):
|
||||
|
@ -19,6 +19,7 @@ class ConfigurationExport(flask_restful.Resource):
|
|||
config_export = plaintext_config
|
||||
if should_encrypt:
|
||||
password = data["password"]
|
||||
config_export = encrypt_config(plaintext_config, password)
|
||||
plaintext_config = json.dumps(plaintext_config)
|
||||
config_export = encrypt_string(plaintext_config, password)
|
||||
|
||||
return {"config_export": config_export, "encrypted": should_encrypt}
|
||||
|
|
|
@ -9,9 +9,10 @@ from flask import request
|
|||
from common.utils.exceptions import InvalidConfigurationError
|
||||
from monkey_island.cc.resources.auth.auth import jwt_required
|
||||
from monkey_island.cc.services.config import ConfigService
|
||||
from monkey_island.cc.services.utils.config_encryption import (
|
||||
from monkey_island.cc.services.utils.encryption import (
|
||||
InvalidCiphertextError,
|
||||
InvalidCredentialsError,
|
||||
decrypt_config,
|
||||
decrypt_ciphertext,
|
||||
is_encrypted,
|
||||
)
|
||||
|
||||
|
@ -68,13 +69,13 @@ class ConfigurationImport(flask_restful.Resource):
|
|||
|
||||
@staticmethod
|
||||
def _get_plaintext_config_from_request(request_contents: dict) -> dict:
|
||||
if ConfigurationImport.is_config_encrypted(request_contents["config"]):
|
||||
return decrypt_config(request_contents["config"], request_contents["password"])
|
||||
else:
|
||||
try:
|
||||
return json.loads(request_contents["config"])
|
||||
except JSONDecodeError:
|
||||
raise InvalidConfigurationError
|
||||
try:
|
||||
config = request_contents["config"]
|
||||
if ConfigurationImport.is_config_encrypted(request_contents["config"]):
|
||||
config = decrypt_ciphertext(config, request_contents["password"])
|
||||
return json.loads(config)
|
||||
except (JSONDecodeError, InvalidCiphertextError):
|
||||
raise InvalidConfigurationError
|
||||
|
||||
@staticmethod
|
||||
def import_config(config_json):
|
||||
|
|
|
@ -1,62 +0,0 @@
|
|||
import base64
|
||||
import io
|
||||
import json
|
||||
import logging
|
||||
from typing import Dict
|
||||
|
||||
import pyAesCrypt
|
||||
|
||||
from common.utils.exceptions import InvalidConfigurationError
|
||||
|
||||
BUFFER_SIZE = pyAesCrypt.crypto.bufferSizeDef
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def encrypt_config(config: Dict, password: str) -> str:
|
||||
plaintext_config_stream = io.BytesIO(json.dumps(config).encode())
|
||||
ciphertext_config_stream = io.BytesIO()
|
||||
|
||||
pyAesCrypt.encryptStream(
|
||||
plaintext_config_stream, ciphertext_config_stream, password, BUFFER_SIZE
|
||||
)
|
||||
|
||||
ciphertext_b64 = base64.b64encode(ciphertext_config_stream.getvalue())
|
||||
logger.info("Configuration encrypted.")
|
||||
|
||||
return ciphertext_b64.decode()
|
||||
|
||||
|
||||
def decrypt_config(ciphertext: str, password: str) -> Dict:
|
||||
ciphertext = base64.b64decode(ciphertext)
|
||||
ciphertext_config_stream = io.BytesIO(ciphertext)
|
||||
dec_plaintext_config_stream = io.BytesIO()
|
||||
|
||||
len_ciphertext_config_stream = len(ciphertext_config_stream.getvalue())
|
||||
|
||||
try:
|
||||
pyAesCrypt.decryptStream(
|
||||
ciphertext_config_stream,
|
||||
dec_plaintext_config_stream,
|
||||
password,
|
||||
BUFFER_SIZE,
|
||||
len_ciphertext_config_stream,
|
||||
)
|
||||
except ValueError as ex:
|
||||
if str(ex).startswith("Wrong password"):
|
||||
logger.info("Wrong password for configuration provided.")
|
||||
raise InvalidCredentialsError
|
||||
else:
|
||||
logger.info("The provided configuration file is corrupt.")
|
||||
raise InvalidConfigurationError
|
||||
plaintext_config = json.loads(dec_plaintext_config_stream.getvalue().decode("utf-8"))
|
||||
return plaintext_config
|
||||
|
||||
|
||||
def is_encrypted(ciphertext: str) -> bool:
|
||||
ciphertext = base64.b64decode(ciphertext)
|
||||
return ciphertext.startswith(b"AES")
|
||||
|
||||
|
||||
class InvalidCredentialsError(Exception):
|
||||
""" Raise when credentials supplied are invalid """
|
|
@ -0,0 +1,59 @@
|
|||
import base64
|
||||
import io
|
||||
import logging
|
||||
|
||||
import pyAesCrypt
|
||||
|
||||
BUFFER_SIZE = pyAesCrypt.crypto.bufferSizeDef
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def encrypt_string(plaintext: str, password: str) -> str:
|
||||
plaintext_stream = io.BytesIO(plaintext.encode())
|
||||
ciphertext_stream = io.BytesIO()
|
||||
|
||||
pyAesCrypt.encryptStream(plaintext_stream, ciphertext_stream, password, BUFFER_SIZE)
|
||||
|
||||
ciphertext_b64 = base64.b64encode(ciphertext_stream.getvalue())
|
||||
logger.info("String encrypted.")
|
||||
|
||||
return ciphertext_b64.decode()
|
||||
|
||||
|
||||
def decrypt_ciphertext(ciphertext: str, password: str) -> str:
|
||||
ciphertext = base64.b64decode(ciphertext)
|
||||
ciphertext_stream = io.BytesIO(ciphertext)
|
||||
plaintext_stream = io.BytesIO()
|
||||
|
||||
ciphertext_stream_len = len(ciphertext_stream.getvalue())
|
||||
|
||||
try:
|
||||
pyAesCrypt.decryptStream(
|
||||
ciphertext_stream,
|
||||
plaintext_stream,
|
||||
password,
|
||||
BUFFER_SIZE,
|
||||
ciphertext_stream_len,
|
||||
)
|
||||
except ValueError as ex:
|
||||
if str(ex).startswith("Wrong password"):
|
||||
logger.info("Wrong password provided for decryption.")
|
||||
raise InvalidCredentialsError
|
||||
else:
|
||||
logger.info("The corrupt ciphertext provided.")
|
||||
raise InvalidCiphertextError
|
||||
return plaintext_stream.getvalue().decode("utf-8")
|
||||
|
||||
|
||||
def is_encrypted(ciphertext: str) -> bool:
|
||||
ciphertext = base64.b64decode(ciphertext)
|
||||
return ciphertext.startswith(b"AES")
|
||||
|
||||
|
||||
class InvalidCredentialsError(Exception):
|
||||
""" Raised when password for decryption is invalid """
|
||||
|
||||
|
||||
class InvalidCiphertextError(Exception):
|
||||
""" Raised when ciphertext is corrupted """
|
|
@ -18,3 +18,8 @@ def monkey_config(data_for_tests_dir):
|
|||
)
|
||||
plaintext_config = json.loads(open(plaintext_monkey_config_standard_path, "r").read())
|
||||
return plaintext_config
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def monkey_config_json(monkey_config):
|
||||
return json.dumps(monkey_config)
|
||||
|
|
|
@ -1,5 +1,3 @@
|
|||
import json
|
||||
|
||||
import pytest
|
||||
from tests.unit_tests.monkey_island.cc.services.utils.ciphertexts_for_encryption_test import (
|
||||
MALFORMED_CIPHER_TEXT_CORRUPTED,
|
||||
|
@ -8,16 +6,15 @@ from tests.unit_tests.monkey_island.cc.services.utils.test_config_encryption imp
|
|||
|
||||
from common.utils.exceptions import InvalidConfigurationError
|
||||
from monkey_island.cc.resources.configuration_import import ConfigurationImport
|
||||
from monkey_island.cc.services.utils.config_encryption import encrypt_config
|
||||
from monkey_island.cc.services.utils.encryption import encrypt_string
|
||||
|
||||
|
||||
def test_is_config_encrypted__json(monkey_config):
|
||||
monkey_config = json.dumps(monkey_config)
|
||||
assert not ConfigurationImport.is_config_encrypted(monkey_config)
|
||||
def test_is_config_encrypted__json(monkey_config_json):
|
||||
assert not ConfigurationImport.is_config_encrypted(monkey_config_json)
|
||||
|
||||
|
||||
def test_is_config_encrypted__ciphertext(monkey_config):
|
||||
encrypted_config = encrypt_config(monkey_config, PASSWORD)
|
||||
def test_is_config_encrypted__ciphertext(monkey_config_json):
|
||||
encrypted_config = encrypt_string(monkey_config_json, PASSWORD)
|
||||
assert ConfigurationImport.is_config_encrypted(encrypted_config)
|
||||
|
||||
|
||||
|
|
|
@ -3,10 +3,10 @@ from tests.unit_tests.monkey_island.cc.services.utils.ciphertexts_for_encryption
|
|||
MALFORMED_CIPHER_TEXT_CORRUPTED,
|
||||
)
|
||||
|
||||
from monkey_island.cc.services.utils.config_encryption import (
|
||||
from monkey_island.cc.services.utils.encryption import (
|
||||
InvalidCredentialsError,
|
||||
decrypt_config,
|
||||
encrypt_config,
|
||||
decrypt_ciphertext,
|
||||
encrypt_string,
|
||||
)
|
||||
|
||||
MONKEY_CONFIGS_DIR_PATH = "monkey_configs"
|
||||
|
@ -15,23 +15,23 @@ PASSWORD = "hello123"
|
|||
INCORRECT_PASSWORD = "goodbye321"
|
||||
|
||||
|
||||
def test_encrypt_decrypt_config(monkey_config):
|
||||
encrypted_config = encrypt_config(monkey_config, PASSWORD)
|
||||
assert decrypt_config(encrypted_config, PASSWORD) == monkey_config
|
||||
def test_encrypt_decrypt_string(monkey_config_json):
|
||||
encrypted_config = encrypt_string(monkey_config_json, PASSWORD)
|
||||
assert decrypt_ciphertext(encrypted_config, PASSWORD) == monkey_config_json
|
||||
|
||||
|
||||
def test_encrypt_decrypt_config__wrong_password(monkey_config):
|
||||
encrypted_config = encrypt_config(monkey_config, PASSWORD)
|
||||
def test_encrypt_decrypt_string__wrong_password(monkey_config_json):
|
||||
encrypted_config = encrypt_string(monkey_config_json, PASSWORD)
|
||||
with pytest.raises(InvalidCredentialsError):
|
||||
decrypt_config(encrypted_config, INCORRECT_PASSWORD)
|
||||
decrypt_ciphertext(encrypted_config, INCORRECT_PASSWORD)
|
||||
|
||||
|
||||
def test_encrypt_decrypt_config__malformed_corrupted():
|
||||
def test_encrypt_decrypt_string__malformed_corrupted():
|
||||
with pytest.raises(ValueError):
|
||||
decrypt_config(MALFORMED_CIPHER_TEXT_CORRUPTED, PASSWORD)
|
||||
decrypt_ciphertext(MALFORMED_CIPHER_TEXT_CORRUPTED, PASSWORD)
|
||||
|
||||
|
||||
def test_encrypt_decrypt_config__decrypt_no_password(monkey_config):
|
||||
encrypted_config = encrypt_config(monkey_config, PASSWORD)
|
||||
def test_encrypt_decrypt_string__decrypt_no_password(monkey_config_json):
|
||||
encrypted_config = encrypt_string(monkey_config_json, PASSWORD)
|
||||
with pytest.raises(InvalidCredentialsError):
|
||||
decrypt_config(encrypted_config, "")
|
||||
decrypt_ciphertext(encrypted_config, "")
|
||||
|
|
Loading…
Reference in New Issue