Add exception handling for config decryption

This commit is contained in:
Shreya 2021-06-01 15:45:20 +05:30 committed by VakarisZ
parent d67e84a6a7
commit b9fb4c6902
3 changed files with 40 additions and 19 deletions

View File

@ -64,3 +64,7 @@ class NoCredentialsError(Exception):
class InvalidConfigurationError(Exception): class InvalidConfigurationError(Exception):
""" Raise when configuration is invalid """ """ Raise when configuration is invalid """
class FailedDecryption(Exception):
""" Raise when any kind of decryption fails """

View File

@ -6,10 +6,12 @@ from flask import request
from common.utils.exceptions import ( from common.utils.exceptions import (
InvalidConfigurationError, InvalidConfigurationError,
InvalidCredentialsError, # InvalidCredentialsError,
NoCredentialsError, NoCredentialsError,
FailedDecryption,
) )
from monkey_island.cc.resources.auth.auth import jwt_required from monkey_island.cc.resources.auth.auth import jwt_required
from monkey_island.cc.services.utils.config_encryption import decrypt_config
@dataclass @dataclass
@ -30,12 +32,19 @@ class TempConfiguration(flask_restful.Resource):
def post(self): def post(self):
request_contents = json.loads(request.data) request_contents = json.loads(request.data)
try: try:
self.decrypt(request_contents["password"]) decrypt_config(request_contents["encrypted_config"], request_contents["password"])
self.import_config() self.import_config()
return ResponseContents().form_response() return ResponseContents().form_response()
except InvalidCredentialsError: # except InvalidCredentialsError:
# return ResponseContents(
# import_status="wrong_password", message="Wrong password supplied", status_code=403
# ).form_response()
except FailedDecryption as ex:
return ResponseContents( return ResponseContents(
import_status="wrong_password", message="Wrong password supplied", status_code=403 import_status="decryption_failure",
message="Decryptioon of configuration failed. Error thrown during decryption: "
+ f"{str(ex)}",
status_code=403,
).form_response() ).form_response()
except InvalidConfigurationError: except InvalidConfigurationError:
return ResponseContents( return ResponseContents(
@ -52,12 +61,12 @@ class TempConfiguration(flask_restful.Resource):
status_code=403, status_code=403,
).form_response() ).form_response()
def decrypt(self, password=""): # def decrypt(self, password=""):
if not password: # if not password:
raise NoCredentialsError # raise NoCredentialsError
if not password == "abc": # if not password == "abc":
raise InvalidCredentialsError # raise InvalidCredentialsError
return False # return False
def import_config(self): def import_config(self):
return True return True

View File

@ -4,6 +4,8 @@ from typing import Dict
import pyAesCrypt import pyAesCrypt
from common.utils.exceptions import FailedDecryption, NoCredentialsError
# TODO use from pyAesCrypt # TODO use from pyAesCrypt
BUFFER_SIZE = 64 * 1024 BUFFER_SIZE = 64 * 1024
@ -21,11 +23,15 @@ def encrypt_config(config: Dict, password: str) -> str:
def decrypt_config(enc_config: bytes, password: str) -> Dict: def decrypt_config(enc_config: bytes, password: str) -> Dict:
if not password:
raise NoCredentialsError
ciphertext_config_stream = io.BytesIO(enc_config) ciphertext_config_stream = io.BytesIO(enc_config)
dec_plaintext_config_stream = io.BytesIO() dec_plaintext_config_stream = io.BytesIO()
len_ciphertext_config_stream = len(ciphertext_config_stream.getvalue()) len_ciphertext_config_stream = len(ciphertext_config_stream.getvalue())
try:
pyAesCrypt.decryptStream( pyAesCrypt.decryptStream(
ciphertext_config_stream, ciphertext_config_stream,
dec_plaintext_config_stream, dec_plaintext_config_stream,
@ -33,6 +39,8 @@ def decrypt_config(enc_config: bytes, password: str) -> Dict:
BUFFER_SIZE, BUFFER_SIZE,
len_ciphertext_config_stream, len_ciphertext_config_stream,
) )
except ValueError as ex:
raise FailedDecryption(str(ex))
plaintext_config = json.loads(dec_plaintext_config_stream.getvalue().decode("utf-8")) plaintext_config = json.loads(dec_plaintext_config_stream.getvalue().decode("utf-8"))
return plaintext_config return plaintext_config