address review

lean/commands/cloud/push.py

@@ -19,7 +19,6 @@
 from lean.click import LeanCommand, PathParameter
 from lean.constants import PROJECT_CONFIG_FILE_NAME
 from lean.container import container
-from lean.components.util.encryption_helper import get_project_key_hash
 from lean.models.encryption import ActionType
 
 @command(cls=LeanCommand)
@@ -64,6 +63,7 @@ def push(project: Optional[Path], encrypt: Optional[bool], decrypt: Optional[boo
             raise RuntimeError(f"'{project}' is not a Lean project")
 
         if encrypt and key is not None:
+            from lean.components.util.encryption_helper import get_project_key_hash
             # lets check if the given key is registered with the cloud
             organization_id = container.organization_manager.try_get_working_organization_id()
             available_encryption_keys = container.api_client.encryption_keys.list(organization_id)['keys']

lean/commands/decrypt.py

@@ -17,7 +17,7 @@
 
 from lean.click import LeanCommand, PathParameter
 from lean.container import container
-from lean.components.util.encryption_helper import get_decrypted_file_content_for_project
+
 
 @command(cls=LeanCommand)
 @argument("project", type=PathParameter(exists=True, file_okay=False, dir_okay=True))
@@ -38,24 +38,19 @@ def decrypt(project: Path,
 
     # Check if the project is already decrypted
     if not project_config.get("encrypted", False):
+        logger.info(f"Successfully decrypted project {project}")
         return
 
     decryption_key: Path = project_config.get('encryption-key-path', None)
-    if decryption_key is not None and Path(decryption_key).exists():
-        decryption_key = Path(decryption_key)
-
-    if decryption_key is None and key is None:
-        raise RuntimeError("No decryption key was provided, please provide one using --key")
-    elif decryption_key is None:
-        decryption_key = key
-    elif key is not None and decryption_key != key:
-        raise RuntimeError(f"Provided decryption key ({key}) does not match the decryption key in the project ({decryption_key})")
+    from lean.components.util.encryption_helper import get_and_validate_user_input_encryption_key
+    decryption_key = get_and_validate_user_input_encryption_key(key, decryption_key)
 
     organization_id = container.organization_manager.try_get_working_organization_id()
 
     source_files = project_manager.get_source_files(project)
     try:
-        decrypted_data = get_decrypted_file_content_for_project(project, 
+        from lean.components.util.encryption_helper import get_decrypted_file_content_for_local_project
+        decrypted_data = get_decrypted_file_content_for_local_project(project, 
                                 source_files, decryption_key, project_config_manager, organization_id)
     except Exception as e:
         raise RuntimeError(f"Could not decrypt project {project}: {e}")

lean/commands/encrypt.py

@@ -17,7 +17,6 @@
 
 from lean.click import LeanCommand, PathParameter
 from lean.container import container
-from lean.components.util.encryption_helper import get_encrypted_file_content_for_project
 
 
 @command(cls=LeanCommand)
@@ -40,24 +39,19 @@ def encrypt(project: Path,
 
     # Check if the project is already encrypted
     if project_config.get('encrypted', False):
+        logger.info(f"Local files encrypted successfully.")
         return
 
     encryption_key: Path = project_config.get('encryption-key-path', None)
-    if encryption_key is not None and Path(encryption_key).exists():
-        encryption_key = Path(encryption_key)
-
-    if encryption_key is None and key is None:
-        raise RuntimeError("No encryption key was provided, please provide one using --key")
-    elif encryption_key is None:
-        encryption_key = key
-    elif key is not None and encryption_key != key:
-        raise RuntimeError(f"Provided encryption key ({key}) does not match the encryption key in the project ({encryption_key})")
+    from lean.components.util.encryption_helper import get_and_validate_user_input_encryption_key
+    encryption_key = get_and_validate_user_input_encryption_key(key, encryption_key)
 
     organization_id = container.organization_manager.try_get_working_organization_id()
 
     source_files = project_manager.get_source_files(project)
     try:
-        encrypted_data = get_encrypted_file_content_for_project(project, 
+        from lean.components.util.encryption_helper import get_encrypted_file_content_for_local_project
+        encrypted_data = get_encrypted_file_content_for_local_project(project, 
                             source_files, encryption_key, project_config_manager, organization_id)
     except Exception as e:
         raise RuntimeError(f"Could not encrypt project {project}: {e}")

lean/components/cloud/pull_manager.py

@@ -25,8 +25,6 @@
 from lean.models.utils import LeanLibraryReference
 from lean.models.encryption import ActionType
 from lean.components.config.storage import safe_save
-from lean.components.util.encryption_helper import get_decrypted_content_from_cloud_project, \
-                                    get_encrypted_content_from_cloud_project, get_project_key_hash
 
 class PullManager:
     """The PullManager class is responsible for synchronizing cloud projects to the local drive."""
@@ -182,10 +180,8 @@ def _pull_project(self, project: QCProject, encryption_action: Optional[ActionTy
         local_encryption_state = project_config.get("encrypted", False)
 
         # Handle mismatch cases
-        if not encryption_key and bool(project.encrypted) != bool(local_encryption_state):
-            raise RuntimeError(f"Project encryption state mismatch. Please provide encryption key to pull project '{project.name}'")
-        if getattr(project.encryptionKey, 'id', None) and encryption_key and get_project_key_hash(encryption_key) != project.encryptionKey.id:
-            raise RuntimeError(f"Encryption Key mismatch. Please provide correct encryption key to pull project '{project.name}'")
+        from lean.components.util.encryption_helper import validate_key_and_encryption_state_for_cloud_project
+        validate_key_and_encryption_state_for_cloud_project(project, local_encryption_state, encryption_key)
 
         # Pull the cloud files to the local drive
         self._pull_files(project, local_project_path, encryption_action, encryption_key)
@@ -198,9 +194,10 @@ def _pull_project(self, project: QCProject, encryption_action: Optional[ActionTy
         project_config.set("description", project.description)
         project_config.set("organization-id", project.organizationId)
         project_config.set("python-venv", project.leanEnvironment)
-        project_config.set('encrypted', project.encrypted)
         if encryption_key:
             project_config.set("encrypted", encryption_action == ActionType.ENCRYPT)
+        else:
+            project_config.set('encrypted', project.encrypted)
 
         if not project.leanPinnedToMaster:
             project_config.set("lean-engine", project.leanVersionId)
@@ -220,11 +217,9 @@ def _pull_files(self, project: QCProject, local_project_path: Path, encryption_a
 
         cloud_files = self._api_client.files.get_all(project.projectId)
         if encryption_key:
+            from lean.components.util.encryption_helper import get_appropriate_files_from_cloud_project
             organization_id = self._organization_manager.try_get_working_organization_id()
-            if encryption_action == ActionType.DECRYPT:
-                cloud_files = get_decrypted_content_from_cloud_project(project, cloud_files, encryption_key, organization_id)
-            if encryption_action == ActionType.ENCRYPT:
-                cloud_files = get_encrypted_content_from_cloud_project(project, cloud_files, encryption_key, organization_id)
+            cloud_files = get_appropriate_files_from_cloud_project(project, cloud_files, encryption_key, organization_id, encryption_action)
 
         for cloud_file in cloud_files:
             self._last_file = cloud_file.name

lean/components/cloud/push_manager.py

@@ -21,8 +21,6 @@
 from lean.components.util.project_manager import ProjectManager
 from lean.models.api import QCLanguage, QCProject
 from lean.models.utils import LeanLibraryReference
-from lean.components.util.encryption_helper import get_encrypted_file_content_for_project,\
-                                get_project_key_hash, get_decrypted_file_content_for_project
 from lean.models.encryption import ActionType
 
 class PushManager:
@@ -150,10 +148,8 @@ def _push_project(self, project_path: Path, organization_id: str, encryption_act
             self._logger.info(f"Successfully created cloud project '{cloud_project.name}'{organization_message_part}")
 
         # Handle mismatch cases
-        if not encryption_key and bool(cloud_project.encrypted) != bool(local_encryption_state):
-            raise RuntimeError(f"Project encryption state mismatch. Please provide encryption key to push the project.")
-        if getattr(cloud_project.encryptionKey, 'id', None) and encryption_key and get_project_key_hash(encryption_key) != cloud_project.encryptionKey.id:
-            raise RuntimeError(f"Project encryption key mismatch. Please provide correct encryption key to push the project.")
+        from lean.components.util.encryption_helper import validate_key_and_encryption_state_for_cloud_project
+        validate_key_and_encryption_state_for_cloud_project(cloud_project, local_encryption_state, encryption_key)
 
         # Finalize pushing by updating locally modified metadata, files and libraries
         self._push_metadata(project_path, cloud_project, encryption_action, encryption_key)
@@ -164,14 +160,10 @@ def _get_files(self, project: Path, encryption_action: Optional[ActionType], enc
         :param project: the local project to push the files of
         """
         paths = self._project_manager.get_source_files(project)
-        if (encryption_key):
+        if encryption_key:
+            from lean.components.util.encryption_helper import get_appropriate_files_from_local_project
             organization_id = self._organization_manager.try_get_working_organization_id()
-            if encryption_action == ActionType.ENCRYPT:
-                data = get_encrypted_file_content_for_project(project, paths,
-                        encryption_key, self._project_config_manager, organization_id)
-            else:
-                data = get_decrypted_file_content_for_project(project,
-                        paths, encryption_key, self._project_config_manager, organization_id)
+            data = get_appropriate_files_from_local_project(project, paths, encryption_key, self._project_config_manager, organization_id, encryption_action)
             files = [
             {
                 'name': path.relative_to(project).as_posix(),
@@ -243,6 +235,7 @@ def _push_metadata(self, project: Path, cloud_project: QCProject, encryption_act
 
         if (encryption_action is not None and encryption_key is not None):
             if encryption_action == ActionType.ENCRYPT:
+                from lean.components.util.encryption_helper import get_project_key_hash
                 encryption_key_id = get_project_key_hash(encryption_key)
                 update_args["encryption_key"] = encryption_key_id
             else: