bug: keep the key path and leave libraries alone on cloud pull --encrypt

lean/components/cloud/pull_manager.py

@@ -101,6 +101,9 @@ def pull_projects(self, projects_to_pull: List[QCProject], all_cloud_projects: O
         :param projects_to_pull: the cloud projects that need to be pulled
         :param all_cloud_projects: all the projects available in the cloud
         """
+        # libraries are pulled as they are in the cloud, only the given projects are encrypted or decrypted
+        project_ids_to_transform = {project.projectId for project in projects_to_pull}
+
         if all_cloud_projects is not None:
             projects, inaccessible_libraries = self._project_manager.get_cloud_projects_libraries(all_cloud_projects,
                                                                                                projects_to_pull)
@@ -137,7 +140,11 @@ def pull_projects(self, projects_to_pull: List[QCProject], all_cloud_projects: O
         for index, project in enumerate(projects_to_pull, start=1):
             try:
                 self._logger.info(f"[{index}/{len(projects_to_pull)}] Pulling '{project.name}'")
-                projects_with_paths.append((project, self._pull_project(project, encryption_action, encryption_key)))
+                if project.projectId in project_ids_to_transform:
+                    local_path = self._pull_project(project, encryption_action, encryption_key)
+                else:
+                    local_path = self._pull_project(project, None, None)
+                projects_with_paths.append((project, local_path))
             except Exception as ex:
                 from traceback import format_exc
                 self._logger.debug(format_exc().strip())
@@ -198,6 +205,10 @@ def _pull_project(self, project: QCProject, encryption_action: Optional[ActionTy
         project_config.set("python-venv", project.leanEnvironment)
         if encryption_key:
             project_config.set("encrypted", encryption_action == ActionType.ENCRYPT)
+            if encryption_action == ActionType.ENCRYPT:
+                project_config.set("encryption-key-path", str(encryption_key))
+            else:
+                project_config.delete("encryption-key-path")
         else:
             project_config.set('encrypted', project.encrypted)
 

tests/commands/cloud/test_pull.py

@@ -15,7 +15,7 @@
 from pathlib import Path
 from datetime import datetime
 from click.testing import CliRunner
-from lean.models.api import QCFullFile
+from lean.models.api import QCFullFile, QCProjectLibrary
 from lean.commands import lean
 from lean.container import container
 from tests.conftest import initialize_container
@@ -506,3 +506,106 @@ def _get_expected_encrypted_files_content() -> dict:
 UGw0ehtO8qY5FmPGcUlkBGuqmd7r6aLE4mosoZrc/UyZb+clWNYJITRLFJbQpWm3EU/Xrt5UM8uWwEdV
 bFWAAkX56MyDHwJefC1nkA=="""
 }
+
+def test_cloud_pull_stores_the_key_path_with_encrypted_flag_given() -> None:
+    create_fake_lean_cli_directory()
+
+    project_path = Path.cwd() / "Python Project"
+    encryption_file_path = project_path / "encryption.txt"
+    encryption_file_path.write_text("KtSwJtq5a4uuQmxbPqcCP3d8yMRz5TZxDBAKy7kGwPcvcvsNBdCprGYwSBN8ntJa5JNNYHTB2GrBpAbkA38kCdnceegffZH7")
+
+    cloud_project = create_api_project(1, "Python Project")
+
+    api_client = mock.Mock()
+    api_client.projects.get_all.return_value = [cloud_project]
+    api_client.projects.get.return_value = cloud_project
+    api_client.files.get_all = mock.MagicMock(return_value=[
+        QCFullFile(name="main.py", content="print(1)\n", modified=datetime.now(), isLibrary=False)
+    ])
+
+    init_container(api_client_to_use=api_client)
+    project_config = container.project_config_manager.get_project_config(project_path)
+    project_config.set("cloud-id", 1)
+
+    result = CliRunner().invoke(lean, ["cloud", "pull", "--project", project_path.name, "--encrypt", "--key", encryption_file_path])
+
+    assert result.exit_code == 0
+
+    project_config = container.project_config_manager.get_project_config(project_path)
+    assert project_config.get("encrypted", False) == True
+    assert project_config.get("encryption-key-path", None) == str(encryption_file_path)
+
+    # a later push without --key must keep the project encrypted in the cloud
+    key_hash = get_project_key_hash(encryption_file_path)
+    api_client.encryption_keys.list = mock.MagicMock(return_value={'keys': [{'name': 'test', 'hash': key_hash}]})
+    api_client.projects.update.reset_mock()
+    result = CliRunner().invoke(lean, ["cloud", "push", "--project", project_path.name])
+
+    assert result.exit_code == 0
+    assert api_client.projects.update.call_args.kwargs["encryption_key"] == key_hash
+
+def test_cloud_pull_removes_the_key_path_with_decrypted_flag_given() -> None:
+    create_fake_lean_cli_directory()
+
+    project_path = Path.cwd() / "Python Project"
+    encryption_file_path = project_path / "encryption.txt"
+    encryption_file_path.write_text("KtSwJtq5a4uuQmxbPqcCP3d8yMRz5TZxDBAKy7kGwPcvcvsNBdCprGYwSBN8ntJa5JNNYHTB2GrBpAbkA38kCdnceegffZH7")
+    key_hash = get_project_key_hash(encryption_file_path)
+
+    cloud_project = create_api_project(1, "Python Project", encrypted=True, encryptionKey={"name":"test", "id": key_hash})
+
+    api_client = mock.Mock()
+    api_client.projects.get_all.return_value = [cloud_project]
+    api_client.projects.get.return_value = cloud_project
+    fake_cloud_files = [QCFullFile(name=name, content=content, modified=datetime.now(), isLibrary=False)
+                        for name, content in _get_expected_encrypted_files_content().items()]
+    api_client.files.get_all = mock.MagicMock(return_value=fake_cloud_files)
+
+    init_container(api_client_to_use=api_client)
+    project_config = container.project_config_manager.get_project_config(project_path)
+    project_config.set("encrypted", True)
+    project_config.set("encryption-key-path", str(encryption_file_path))
+    project_config.set("cloud-id", 1)
+
+    result = CliRunner().invoke(lean, ["cloud", "pull", "--project", project_path.name, "--decrypt", "--key", encryption_file_path])
+
+    assert result.exit_code == 0
+
+    project_config = container.project_config_manager.get_project_config(project_path)
+    assert project_config.get("encrypted", False) == False
+    assert project_config.get("encryption-key-path", None) is None
+
+def test_cloud_pull_does_not_encrypt_referenced_libraries_with_encrypted_flag_given() -> None:
+    create_fake_lean_cli_directory()
+
+    project_path = Path.cwd() / "Python Project"
+    encryption_file_path = project_path / "encryption.txt"
+    encryption_file_path.write_text("KtSwJtq5a4uuQmxbPqcCP3d8yMRz5TZxDBAKy7kGwPcvcvsNBdCprGYwSBN8ntJa5JNNYHTB2GrBpAbkA38kCdnceegffZH7")
+
+    cloud_project = create_api_project(1, "Python Project")
+    library_project = create_api_project(2, "Library/Shared Library")
+    cloud_project.libraries = [QCProjectLibrary(projectId=2, libraryName="Library/Shared Library", ownerName="test", access=True)]
+
+    api_client = mock.Mock()
+    api_client.projects.get_all.return_value = [cloud_project, library_project]
+    api_client.projects.get.side_effect = lambda project_id, organization_id: {1: cloud_project, 2: library_project}[project_id]
+    cloud_files = {
+        1: [QCFullFile(name="main.py", content="print(1)\n", modified=datetime.now(), isLibrary=False)],
+        2: [QCFullFile(name="main.py", content="LIBRARY_VALUE = 1\n", modified=datetime.now(), isLibrary=False)]
+    }
+    api_client.files.get_all.side_effect = lambda project_id: [file.model_copy() for file in cloud_files[project_id]]
+
+    init_container(api_client_to_use=api_client)
+    project_config = container.project_config_manager.get_project_config(project_path)
+    project_config.set("cloud-id", 1)
+
+    result = CliRunner().invoke(lean, ["cloud", "pull", "--project", project_path.name, "--encrypt", "--key", encryption_file_path])
+
+    assert result.exit_code == 0
+
+    library_path = Path.cwd() / "Library" / "Shared Library"
+    library_config = container.project_config_manager.get_project_config(library_path)
+    assert library_config.get("encrypted", False) == False
+    assert library_config.get("encryption-key-path", None) is None
+    assert (library_path / "main.py").read_text() == "LIBRARY_VALUE = 1\n"
+    assert container.project_config_manager.get_project_config(project_path).get("encrypted", False) == True