From b9fb4c6902f668a3698026a26795ef040e8b7e3e Mon Sep 17 00:00:00 2001
From: Shreya <shreya.malviya@gmail.com>
Date: Tue, 1 Jun 2021 15:45:20 +0530
Subject: [PATCH] Add exception handling for config decryption

---
 monkey/common/utils/exceptions.py             |  8 +++--
 .../cc/resources/temp_configuration.py        | 29 ++++++++++++-------
 .../cc/services/utils/config_encryption.py    | 22 +++++++++-----
 3 files changed, 40 insertions(+), 19 deletions(-)

diff --git a/monkey/common/utils/exceptions.py b/monkey/common/utils/exceptions.py
index 0658b74f3..a9b4bd550 100644
--- a/monkey/common/utils/exceptions.py
+++ b/monkey/common/utils/exceptions.py
@@ -55,12 +55,16 @@ class DomainControllerNameFetchError(FailedExploitationError):
 
 
 class InvalidCredentialsError(Exception):
-    """ Raise when credentials supplied are invalid"""
+    """ Raise when credentials supplied are invalid """
 
 
 class NoCredentialsError(Exception):
-    """ Raise when no credentials have been supplied"""
+    """ Raise when no credentials have been supplied """
 
 
 class InvalidConfigurationError(Exception):
     """ Raise when configuration is invalid """
+
+
+class FailedDecryption(Exception):
+    """ Raise when any kind of decryption fails """
diff --git a/monkey/monkey_island/cc/resources/temp_configuration.py b/monkey/monkey_island/cc/resources/temp_configuration.py
index fe9ae09b6..4ec651f63 100644
--- a/monkey/monkey_island/cc/resources/temp_configuration.py
+++ b/monkey/monkey_island/cc/resources/temp_configuration.py
@@ -6,10 +6,12 @@ from flask import request
 
 from common.utils.exceptions import (
     InvalidConfigurationError,
-    InvalidCredentialsError,
+    # InvalidCredentialsError,
     NoCredentialsError,
+    FailedDecryption,
 )
 from monkey_island.cc.resources.auth.auth import jwt_required
+from monkey_island.cc.services.utils.config_encryption import decrypt_config
 
 
 @dataclass
@@ -30,12 +32,19 @@ class TempConfiguration(flask_restful.Resource):
     def post(self):
         request_contents = json.loads(request.data)
         try:
-            self.decrypt(request_contents["password"])
+            decrypt_config(request_contents["encrypted_config"], request_contents["password"])
             self.import_config()
             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(
-                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()
         except InvalidConfigurationError:
             return ResponseContents(
@@ -52,12 +61,12 @@ class TempConfiguration(flask_restful.Resource):
                 status_code=403,
             ).form_response()
 
-    def decrypt(self, password=""):
-        if not password:
-            raise NoCredentialsError
-        if not password == "abc":
-            raise InvalidCredentialsError
-        return False
+    # def decrypt(self, password=""):
+    #     if not password:
+    #         raise NoCredentialsError
+    #     if not password == "abc":
+    #         raise InvalidCredentialsError
+    #     return False
 
     def import_config(self):
         return True
diff --git a/monkey/monkey_island/cc/services/utils/config_encryption.py b/monkey/monkey_island/cc/services/utils/config_encryption.py
index ada10aab2..35b580b78 100644
--- a/monkey/monkey_island/cc/services/utils/config_encryption.py
+++ b/monkey/monkey_island/cc/services/utils/config_encryption.py
@@ -4,6 +4,8 @@ from typing import Dict
 
 import pyAesCrypt
 
+from common.utils.exceptions import FailedDecryption, NoCredentialsError
+
 # TODO use from pyAesCrypt
 BUFFER_SIZE = 64 * 1024
 
@@ -21,18 +23,24 @@ def encrypt_config(config: Dict, password: str) -> str:
 
 
 def decrypt_config(enc_config: bytes, password: str) -> Dict:
+    if not password:
+        raise NoCredentialsError
+
     ciphertext_config_stream = io.BytesIO(enc_config)
     dec_plaintext_config_stream = io.BytesIO()
 
     len_ciphertext_config_stream = len(ciphertext_config_stream.getvalue())
 
-    pyAesCrypt.decryptStream(
-        ciphertext_config_stream,
-        dec_plaintext_config_stream,
-        password,
-        BUFFER_SIZE,
-        len_ciphertext_config_stream,
-    )
+    try:
+        pyAesCrypt.decryptStream(
+            ciphertext_config_stream,
+            dec_plaintext_config_stream,
+            password,
+            BUFFER_SIZE,
+            len_ciphertext_config_stream,
+        )
+    except ValueError as ex:
+        raise FailedDecryption(str(ex))
 
     plaintext_config = json.loads(dec_plaintext_config_stream.getvalue().decode("utf-8"))
     return plaintext_config