Split up the initialization of mongo_key into 2 parts: directory of mongo key initialization that happens during launch and initialization of key which happens after login or registration

monkey/monkey_island/cc/resources/auth/auth.py

@@ -1,4 +1,3 @@
-import json
 import logging
 from functools import wraps
 
@@ -9,8 +8,13 @@ from flask_jwt_extended.exceptions import JWTExtendedException
 from jwt import PyJWTError
 
 import monkey_island.cc.environment.environment_singleton as env_singleton
-import monkey_island.cc.resources.auth.password_utils as password_utils
 import monkey_island.cc.resources.auth.user_store as user_store
+from monkey_island.cc.resources.auth.credential_utils import (
+    get_creds_from_request,
+    get_secret_from_request,
+    password_matches_hash,
+)
+from monkey_island.cc.server_utils.encryption.data_store_encryptor import setup_datastore_key
 
 logger = logging.getLogger(__name__)
 
@@ -38,28 +42,20 @@ class Authenticate(flask_restful.Resource):
             "password": "my_password"
         }
         """
-        (username, password) = _get_credentials_from_request(request)
+        username, password = get_creds_from_request(request)
 
         if _credentials_match_registered_user(username, password):
+            setup_datastore_key(get_secret_from_request(request))
             access_token = _create_access_token(username)
             return make_response({"access_token": access_token, "error": ""}, 200)
         else:
             return make_response({"error": "Invalid credentials"}, 401)
 
 
-def _get_credentials_from_request(request):
-    credentials = json.loads(request.data)
-
-    username = credentials["username"]
-    password = credentials["password"]
-
-    return (username, password)
-
-
-def _credentials_match_registered_user(username, password):
+def _credentials_match_registered_user(username: str, password: str):
     user = user_store.UserStore.username_table.get(username, None)
 
-    if user and password_utils.password_matches_hash(password, user.secret):
+    if user and password_matches_hash(password, user.secret):
         return True
 
     return False

monkey/monkey_island/cc/resources/auth/credential_utils.py

@@ -0,0 +1,37 @@
+import json
+from typing import Tuple
+
+import bcrypt
+from flask import Request, request
+
+from monkey_island.cc.environment.user_creds import UserCreds
+
+
+def hash_password(plaintext_password):
+    salt = bcrypt.gensalt()
+    password_hash = bcrypt.hashpw(plaintext_password.encode("utf-8"), salt)
+
+    return password_hash.decode()
+
+
+def password_matches_hash(plaintext_password, password_hash):
+    return bcrypt.checkpw(plaintext_password.encode("utf-8"), password_hash.encode("utf-8"))
+
+
+def get_user_credentials_from_request(_request) -> UserCreds:
+    username, password = get_creds_from_request(_request)
+    password_hash = hash_password(password)
+
+    return UserCreds(username, password_hash)
+
+
+def get_secret_from_request(_request) -> str:
+    username, password = get_creds_from_request(_request)
+    return f"{username}:{password}"
+
+
+def get_creds_from_request(_request: Request) -> Tuple[str, str]:
+    cred_dict = json.loads(request.data)
+    username = cred_dict.get("username", "")
+    password = cred_dict.get("password", "")
+    return username, password

monkey/monkey_island/cc/resources/auth/password_utils.py

@@ -1,12 +0,0 @@
-import bcrypt
-
-
-def hash_password(plaintext_password):
-    salt = bcrypt.gensalt()
-    password_hash = bcrypt.hashpw(plaintext_password.encode("utf-8"), salt)
-
-    return password_hash.decode()
-
-
-def password_matches_hash(plaintext_password, password_hash):
-    return bcrypt.checkpw(plaintext_password.encode("utf-8"), password_hash.encode("utf-8"))

monkey/monkey_island/cc/resources/auth/registration.py

@@ -1,13 +1,15 @@
-import json
 import logging
 
 import flask_restful
 from flask import make_response, request
 
 import monkey_island.cc.environment.environment_singleton as env_singleton
-import monkey_island.cc.resources.auth.password_utils as password_utils
 from common.utils.exceptions import InvalidRegistrationCredentialsError, RegistrationNotNeededError
-from monkey_island.cc.environment.user_creds import UserCreds
+from monkey_island.cc.resources.auth.credential_utils import (
+    get_secret_from_request,
+    get_user_credentials_from_request,
+)
+from monkey_island.cc.server_utils.encryption.data_store_encryptor import setup_datastore_key
 from monkey_island.cc.setup.mongo.database_initializer import reset_database
 
 logger = logging.getLogger(__name__)
@@ -19,21 +21,13 @@ class Registration(flask_restful.Resource):
         return {"needs_registration": is_registration_needed}
 
     def post(self):
-        credentials = _get_user_credentials_from_request(request)
+        # TODO delete the old key here, before creating new one
+        credentials = get_user_credentials_from_request(request)
 
         try:
             env_singleton.env.try_add_user(credentials)
+            setup_datastore_key(get_secret_from_request(request))
             reset_database()
             return make_response({"error": ""}, 200)
         except (InvalidRegistrationCredentialsError, RegistrationNotNeededError) as e:
             return make_response({"error": str(e)}, 400)
-
-
-def _get_user_credentials_from_request(request):
-    cred_dict = json.loads(request.data)
-
-    username = cred_dict.get("user", "")
-    password = cred_dict.get("password", "")
-    password_hash = password_utils.hash_password(password)
-
-    return UserCreds(username, password_hash)

monkey/monkey_island/cc/server_utils/encryption/data_store_encryptor.py

@@ -1,8 +1,12 @@
+from __future__ import annotations
+
 import io
 import os
 
 # PyCrypto is deprecated, but we use pycryptodome, which uses the exact same imports but
 # is maintained.
+from typing import Union
+
 from Crypto import Random  # noqa: DUO133  # nosec: B413
 
 from monkey_island.cc.server_utils.encryption import KeyBasedEncryptor
@@ -11,37 +15,42 @@ from monkey_island.cc.server_utils.encryption.password_based_byte_encryption imp
 )
 from monkey_island.cc.server_utils.file_utils import open_new_securely_permissioned_file
 
-_encryptor = None
+_encryptor: Union[None, DataStoreEncryptor] = None
 
 
 class DataStoreEncryptor:
     _BLOCK_SIZE = 32
     _KEY_FILENAME = "mongo_key.bin"
 
-    def __init__(self, key_file_dir: str, secret: str):
-        key_file = os.path.join(key_file_dir, self._KEY_FILENAME)
+    def __init__(self, key_file_dir: str):
+        self.key_file_path = os.path.join(key_file_dir, self._KEY_FILENAME)
+        self._key_base_encryptor = None
 
-        if os.path.exists(key_file):
-            self._load_existing_key(key_file, secret)
+    def init_key(self, secret: str):
+        if os.path.exists(self.key_file_path):
+            self._load_existing_key(secret)
         else:
-            self._init_key(key_file, secret)
+            self._create_new_key(secret)
 
-        self._key_base_encryptor = KeyBasedEncryptor(self._cipher_key)
-
-    def _init_key(self, password_file_path: str, secret: str):
-        self._cipher_key = Random.new().read(self._BLOCK_SIZE)
-        encrypted_key = (
-            PasswordBasedByteEncryptor(secret).encrypt(io.BytesIO(self._cipher_key)).getvalue()
-        )
-        with open_new_securely_permissioned_file(password_file_path, "wb") as f:
-            f.write(encrypted_key)
-
-    def _load_existing_key(self, key_file_path: str, secret: str):
-        with open(key_file_path, "rb") as f:
+    def _load_existing_key(self, secret: str):
+        with open(self.key_file_path, "rb") as f:
             encrypted_key = f.read()
-        self._cipher_key = (
+        cipher_key = (
             PasswordBasedByteEncryptor(secret).decrypt(io.BytesIO(encrypted_key)).getvalue()
         )
+        self._key_base_encryptor = KeyBasedEncryptor(cipher_key)
+
+    def _create_new_key(self, secret: str):
+        cipher_key = Random.new().read(self._BLOCK_SIZE)
+        encrypted_key = (
+            PasswordBasedByteEncryptor(secret).encrypt(io.BytesIO(cipher_key)).getvalue()
+        )
+        with open_new_securely_permissioned_file(self.key_file_path, "wb") as f:
+            f.write(encrypted_key)
+        self._key_base_encryptor = KeyBasedEncryptor(cipher_key)
+
+    def is_key_setup(self) -> bool:
+        return self._key_base_encryptor is not None
 
     def enc(self, message: str):
         return self._key_base_encryptor.encrypt(message)
@@ -50,10 +59,22 @@ class DataStoreEncryptor:
         return self._key_base_encryptor.decrypt(enc_message)
 
 
-def initialize_datastore_encryptor(key_file_dir: str, secret: str):
+def initialize_datastore_encryptor(key_file_dir: str):
     global _encryptor
 
-    _encryptor = DataStoreEncryptor(key_file_dir, secret)
+    _encryptor = DataStoreEncryptor(key_file_dir)
+
+
+class EncryptorNotInitializedError(Exception):
+    pass
+
+
+def setup_datastore_key(secret: str):
+    if _encryptor is None:
+        raise EncryptorNotInitializedError
+    else:
+        if not _encryptor.is_key_setup():
+            _encryptor.init_key(secret)
 
 
 def get_datastore_encryptor():

monkey/monkey_island/cc/ui/src/services/AuthService.js

@@ -46,7 +46,7 @@ export default class AuthService {
     return this._authFetch(this.REGISTRATION_API_ENDPOINT, {
       method: 'POST',
       body: JSON.stringify({
-        'user': username,
+        'username': username,
         'password': password
       })
     }).then(res => {

monkey/tests/unit_tests/monkey_island/cc/conftest.py

@@ -11,6 +11,7 @@ from tests.unit_tests.monkey_island.cc.server_utils.encryption.test_password_bas
 )
 
 from monkey_island.cc.server_utils.encryption import initialize_datastore_encryptor
+from monkey_island.cc.server_utils.encryption.data_store_encryptor import setup_datastore_key
 
 
 @pytest.fixture
@@ -32,4 +33,5 @@ ENCRYPTOR_SECRET = "m0nk3y_u53r:53cr3t_p455w0rd"
 
 @pytest.fixture
 def uses_encryptor(data_for_tests_dir):
-    initialize_datastore_encryptor(data_for_tests_dir, ENCRYPTOR_SECRET)
+    initialize_datastore_encryptor(data_for_tests_dir)
+    setup_datastore_key(ENCRYPTOR_SECRET)

monkey/tests/unit_tests/monkey_island/cc/server_utils/encryption/test_data_store_encryptor.py

@@ -8,6 +8,7 @@ from monkey_island.cc.server_utils.encryption import (
     get_datastore_encryptor,
     initialize_datastore_encryptor,
 )
+from monkey_island.cc.server_utils.encryption.data_store_encryptor import setup_datastore_key
 
 PLAINTEXT = "Hello, Monkey!"
 
@@ -22,5 +23,6 @@ def test_encryption(data_for_tests_dir):
 
 
 def test_create_new_password_file(tmpdir):
-    initialize_datastore_encryptor(tmpdir, ENCRYPTOR_SECRET)
+    initialize_datastore_encryptor(tmpdir)
+    setup_datastore_key(ENCRYPTOR_SECRET)
     assert os.path.isfile(os.path.join(tmpdir, DataStoreEncryptor._KEY_FILENAME))