From 88b9e3dc386c1a301e05ccd1ccb35e4de19cb3c6 Mon Sep 17 00:00:00 2001 From: Ronit Jain Date: Tue, 17 Oct 2023 21:43:32 +0530 Subject: [PATCH] Implement Support For Project Encryption (#373) * implement support for project encryption * add remaning tests * fix tests * address review * fix line endings * remove duplicate tests * remove duplicate conversion * add test * self review * fix tests * use debug log instead --- README.md | 42 ++ lean/commands/__init__.py | 4 + lean/commands/cloud/__init__.py | 2 +- lean/commands/cloud/pull.py | 36 +- lean/commands/cloud/push.py | 32 +- lean/commands/decrypt.py | 62 +++ lean/commands/encrypt.py | 64 +++ lean/components/api/api_client.py | 2 + lean/components/api/encryption_keys_client.py | 39 ++ lean/components/api/project_client.py | 6 +- lean/components/cloud/pull_manager.py | 39 +- lean/components/cloud/push_manager.py | 62 ++- lean/components/util/encryption_helper.py | 271 ++++++++++++ lean/container.py | 3 +- lean/models/api.py | 6 + lean/models/encryption.py | 18 + setup.py | 3 +- tests/commands/cloud/test_pull.py | 384 +++++++++++++++++- tests/commands/cloud/test_push.py | 354 +++++++++++++++- tests/commands/test_decrypt.py | 207 ++++++++++ tests/commands/test_encrypt.py | 196 +++++++++ tests/components/util/test_pull_manager.py | 27 +- tests/test_helpers.py | 8 +- 23 files changed, 1803 insertions(+), 64 deletions(-) create mode 100644 lean/commands/decrypt.py create mode 100644 lean/commands/encrypt.py create mode 100644 lean/components/api/encryption_keys_client.py create mode 100644 lean/components/util/encryption_helper.py create mode 100644 lean/models/encryption.py create mode 100644 tests/commands/test_decrypt.py create mode 100644 tests/commands/test_encrypt.py diff --git a/README.md b/README.md index ee54c328..6ca86ff4 100644 --- a/README.md +++ b/README.md @@ -91,7 +91,9 @@ A locally-focused workflow (local development, local execution) with the CLI may - [`lean create-project`](#lean-create-project) - [`lean data download`](#lean-data-download) - [`lean data generate`](#lean-data-generate) +- [`lean decrypt`](#lean-decrypt) - [`lean delete-project`](#lean-delete-project) +- [`lean encrypt`](#lean-encrypt) - [`lean init`](#lean-init) - [`lean library add`](#lean-library-add) - [`lean library remove`](#lean-library-remove) @@ -547,6 +549,9 @@ Usage: lean cloud pull [OPTIONS] Options: --project TEXT Name or id of the project to pull (all cloud projects if not specified) --pull-bootcamp Pull Boot Camp projects (disabled by default) + --encrypt Pull your cloud files and encrypt them before saving on your local drive + --decrypt Pull your cloud files and decrypt them before saving on your local drive + --key FILE Path to the encryption key to use --verbose Enable debug logging --help Show this message and exit. ``` @@ -568,6 +573,9 @@ Usage: lean cloud push [OPTIONS] Options: --project DIRECTORY Path to the local project to push (all local projects if not specified) + --encrypt Push your local files and encrypt them before saving on the cloud + --decrypt Push your local files and decrypt them before saving on the cloud + --key FILE Path to the encryption key to use --verbose Enable debug logging --help Show this message and exit. ``` @@ -797,6 +805,23 @@ Options: _See code: [lean/commands/data/generate.py](lean/commands/data/generate.py)_ +### `lean decrypt` + +Decrypt your local project using the specified decryption key. + +``` +Usage: lean decrypt [OPTIONS] PROJECT + + Decrypt your local project using the specified decryption key. + +Options: + --key FILE Path to the decryption key to use + --verbose Enable debug logging + --help Show this message and exit. +``` + +_See code: [lean/commands/decrypt.py](lean/commands/decrypt.py)_ + ### `lean delete-project` Alias for 'project-delete' @@ -815,6 +840,23 @@ Options: _See code: [lean/commands/delete_project.py](lean/commands/delete_project.py)_ +### `lean encrypt` + +Encrypt your local project using the specified encryption key. + +``` +Usage: lean encrypt [OPTIONS] PROJECT + + Encrypt your local project using the specified encryption key. + +Options: + --key FILE Path to the encryption key to use + --verbose Enable debug logging + --help Show this message and exit. +``` + +_See code: [lean/commands/encrypt.py](lean/commands/encrypt.py)_ + ### `lean init` Scaffold a Lean configuration file and data directory. diff --git a/lean/commands/__init__.py b/lean/commands/__init__.py index c517d8b4..d26e0a4b 100644 --- a/lean/commands/__init__.py +++ b/lean/commands/__init__.py @@ -19,6 +19,8 @@ from lean.commands.create_project import create_project from lean.commands.delete_project import delete_project from lean.commands.data import data +from lean.commands.decrypt import decrypt +from lean.commands.encrypt import encrypt from lean.commands.init import init from lean.commands.library import library from lean.commands.live.live import live @@ -35,6 +37,8 @@ lean.add_command(config) lean.add_command(cloud) lean.add_command(data) +lean.add_command(decrypt) +lean.add_command(encrypt) lean.add_command(library) lean.add_command(live) lean.add_command(login) diff --git a/lean/commands/cloud/__init__.py b/lean/commands/cloud/__init__.py index af5532bf..1f13f325 100644 --- a/lean/commands/cloud/__init__.py +++ b/lean/commands/cloud/__init__.py @@ -35,4 +35,4 @@ def cloud() -> None: cloud.add_command(optimize) cloud.add_command(live) cloud.add_command(status) -cloud.add_command(object_store) +cloud.add_command(object_store) \ No newline at end of file diff --git a/lean/commands/cloud/pull.py b/lean/commands/cloud/pull.py index 348b114c..29e8502a 100644 --- a/lean/commands/cloud/pull.py +++ b/lean/commands/cloud/pull.py @@ -12,23 +12,43 @@ # limitations under the License. from typing import Optional - +from pathlib import Path from click import command, option -from lean.click import LeanCommand +from lean.click import LeanCommand, PathParameter from lean.container import container - +from lean.models.encryption import ActionType @command(cls=LeanCommand) @option("--project", type=str, help="Name or id of the project to pull (all cloud projects if not specified)") @option("--pull-bootcamp", is_flag=True, default=False, help="Pull Boot Camp projects (disabled by default)") -def pull(project: Optional[str], pull_bootcamp: bool) -> None: +@option("--encrypt", + is_flag=True, default=False, + help="Pull your cloud files and encrypt them before saving on your local drive") +@option("--decrypt", + is_flag=True, default=False, + help="Pull your cloud files and decrypt them before saving on your local drive") +@option("--key", + type=PathParameter(exists=True, file_okay=True, dir_okay=False), + help="Path to the encryption key to use") +def pull(project: Optional[str], pull_bootcamp: bool, encrypt: Optional[bool], decrypt: Optional[bool], key: Optional[Path]) -> None: """Pull projects from QuantConnect to the local drive. This command overrides the content of local files with the content of their respective counterparts in the cloud. This command will not delete local files for which there is no counterpart in the cloud. """ + + encryption_action = None + + from lean.components.util.encryption_helper import validate_user_inputs_for_cloud_push_pull_commands + validate_user_inputs_for_cloud_push_pull_commands(encrypt, decrypt, key) + + if encrypt: + encryption_action = ActionType.ENCRYPT + if decrypt: + encryption_action = ActionType.DECRYPT + # Parse which projects need to be pulled project_id = None project_name = None @@ -55,5 +75,11 @@ def pull(project: Optional[str], pull_bootcamp: bool) -> None: if project is None and not pull_bootcamp: projects_to_pull = [p for p in projects_to_pull if not p.name.startswith("Boot Camp/")] + if key is not None and len(projects_to_pull) > 1: + raise RuntimeError(f"Cannot encrypt or decrypt more than one project at a time.") + + # the encryption key info is available when reading the project individually from API + projects_to_pull = [api_client.projects.get(project.projectId, project.organizationId) if project.encrypted == True else project for project in projects_to_pull] + pull_manager = container.pull_manager - pull_manager.pull_projects(projects_to_pull, all_projects) + pull_manager.pull_projects(projects_to_pull, all_projects, encryption_action, key) diff --git a/lean/commands/cloud/push.py b/lean/commands/cloud/push.py index a7a5a233..59168ac2 100644 --- a/lean/commands/cloud/push.py +++ b/lean/commands/cloud/push.py @@ -19,13 +19,22 @@ from lean.click import LeanCommand, PathParameter from lean.constants import PROJECT_CONFIG_FILE_NAME from lean.container import container - +from lean.models.encryption import ActionType @command(cls=LeanCommand) @option("--project", type=PathParameter(exists=True, file_okay=False, dir_okay=True), help="Path to the local project to push (all local projects if not specified)") -def push(project: Optional[Path]) -> None: +@option("--encrypt", + is_flag=True, default=False, + help="Push your local files and encrypt them before saving on the cloud") +@option("--decrypt", + is_flag=True, default=False, + help="Push your local files and decrypt them before saving on the cloud") +@option("--key", + type=PathParameter(exists=True, file_okay=True, dir_okay=False), + help="Path to the encryption key to use") +def push(project: Optional[Path], encrypt: Optional[bool], decrypt: Optional[bool], key: Optional[Path]) -> None: """Push local projects to QuantConnect. This command overrides the content of cloud files with the content of their respective local counterparts. @@ -33,6 +42,15 @@ def push(project: Optional[Path]) -> None: This command will delete cloud files which don't have a local counterpart. """ push_manager = container.push_manager + encryption_action = None + + from lean.components.util.encryption_helper import validate_user_inputs_for_cloud_push_pull_commands + validate_user_inputs_for_cloud_push_pull_commands(encrypt, decrypt, key) + + if encrypt: + encryption_action = ActionType.ENCRYPT + if decrypt: + encryption_action = ActionType.DECRYPT # Parse which projects need to be pushed if project is not None: @@ -41,7 +59,13 @@ def push(project: Optional[Path]) -> None: if not project_config.file.exists(): raise RuntimeError(f"'{project}' is not a Lean project") - push_manager.push_project(project) + if encrypt and key is not None: + from lean.components.util.encryption_helper import validate_encryption_key_registered_with_cloud + validate_encryption_key_registered_with_cloud(key, container.organization_manager, container.api_client) + + push_manager.push_project(project, encryption_action, key) else: + if key is not None: + raise RuntimeError(f"Encryption key can only be specified when pushing a single project.") projects_to_push = [p.parent for p in Path.cwd().rglob(PROJECT_CONFIG_FILE_NAME)] - push_manager.push_projects(projects_to_push) + push_manager.push_projects(projects_to_push, encryption_action, key) diff --git a/lean/commands/decrypt.py b/lean/commands/decrypt.py new file mode 100644 index 00000000..e479b041 --- /dev/null +++ b/lean/commands/decrypt.py @@ -0,0 +1,62 @@ +# QUANTCONNECT.COM - Democratizing Finance, Empowering Individuals. +# Lean CLI v1.0. Copyright 2021 QuantConnect Corporation. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from pathlib import Path +from typing import Optional +from click import command, option, argument + +from lean.click import LeanCommand, PathParameter +from lean.container import container + + +@command(cls=LeanCommand) +@argument("project", type=PathParameter(exists=True, file_okay=False, dir_okay=True)) +@option("--key", + type=PathParameter(exists=True, file_okay=True, dir_okay=False), + help="Path to the decryption key to use") +def decrypt(project: Path, + key: Optional[Path]) -> None: + """Decrypt your local project using the specified decryption key.""" + + logger = container.logger + project_manager = container.project_manager + project_config_manager = container.project_config_manager + project_config = project_config_manager.get_project_config(project) + + # 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) + 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: + 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}") + + for file, decrypted in zip(source_files, decrypted_data): + with open(file, 'w') as f: + f.write(decrypted) + + # Mark the project as decrypted + project_config.set('encrypted', False) + project_config.delete('encryption-key-path') + logger.info(f"Successfully decrypted project {project}") diff --git a/lean/commands/encrypt.py b/lean/commands/encrypt.py new file mode 100644 index 00000000..71e275ef --- /dev/null +++ b/lean/commands/encrypt.py @@ -0,0 +1,64 @@ +# QUANTCONNECT.COM - Democratizing Finance, Empowering Individuals. +# Lean CLI v1.0. Copyright 2021 QuantConnect Corporation. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from pathlib import Path +from typing import Optional +from click import command, option, argument + +from lean.click import LeanCommand, PathParameter +from lean.container import container + + +@command(cls=LeanCommand) +@argument("project", type=PathParameter(exists=True, file_okay=False, dir_okay=True)) +@option("--key", + type=PathParameter(exists=True, file_okay=True, dir_okay=False), + help="Path to the encryption key to use") +def encrypt(project: Path, + key: Optional[Path]) -> None: + """Encrypt your local project using the specified encryption key.""" + + logger = container.logger + project_manager = container.project_manager + project_config_manager = container.project_config_manager + project_config = project_config_manager.get_project_config(project) + + # 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) + 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: + 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}") + for file, encrypted in zip(source_files, encrypted_data): + with open(file, 'w') as f: + f.write(encrypted) + + # Mark the project as encrypted + project_config.set('encrypted', True) + project_config.set('encryption-key-path', str(encryption_key)) + logger.info(f"Local files encrypted successfully with key {encryption_key}") + + + diff --git a/lean/components/api/api_client.py b/lean/components/api/api_client.py index 73a5a1c4..162ee573 100644 --- a/lean/components/api/api_client.py +++ b/lean/components/api/api_client.py @@ -17,6 +17,7 @@ from lean.components.api.backtest_client import BacktestClient from lean.components.api.compile_client import CompileClient from lean.components.api.data_client import DataClient +from lean.components.api.encryption_keys_client import EncryptionKeysClient from lean.components.api.file_client import FileClient from lean.components.api.lean_client import LeanClient from lean.components.api.live_client import LiveClient @@ -55,6 +56,7 @@ def __init__(self, logger: Logger, http_client: HTTPClient, user_id: str, api_to self.backtests = BacktestClient(self) self.compiles = CompileClient(self) self.data = DataClient(self, http_client) + self.encryption_keys = EncryptionKeysClient(self) self.files = FileClient(self) self.live = LiveClient(self) self.market = MarketClient(self) diff --git a/lean/components/api/encryption_keys_client.py b/lean/components/api/encryption_keys_client.py new file mode 100644 index 00000000..ece26a32 --- /dev/null +++ b/lean/components/api/encryption_keys_client.py @@ -0,0 +1,39 @@ +# QUANTCONNECT.COM - Democratizing Finance, Empowering Individuals. +# Lean CLI v1.0. Copyright 2021 QuantConnect Corporation. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from lean.components.api.api_client import * + + +class EncryptionKeysClient: + """The EncryptionKeysClient class contains methods to interact with organizations/encryption/keys/* API endpoints.""" + + def __init__(self, api_client: 'APIClient') -> None: + """Creates a new EncryptionKeysClient instance. + + :param api_client: the APIClient instance to use when making requests + """ + self._api = api_client + + def list(self, organization_id: str) -> str: + """List all values for the given organization + + :param organization_id: the id of the organization who's object store to retrieve from + :return: all objects for the given root key + """ + payload = { + "organizationId": organization_id + } + + data = self._api.post("organizations/encryption/keys/list/", payload) + + return data \ No newline at end of file diff --git a/lean/components/api/project_client.py b/lean/components/api/project_client.py index 2c28617a..b4cbbf70 100644 --- a/lean/components/api/project_client.py +++ b/lean/components/api/project_client.py @@ -82,7 +82,8 @@ def update(self, lean_engine: Optional[int] = None, python_venv: Optional[int] = None, files: Optional[List[Dict[str, str]]] = None, - libraries: Optional[List[int]] = None) -> None: + libraries: Optional[List[int]] = None, + encryption_key: Optional[str] = None) -> None: """Updates an existing project. :param project_id: the id of the project to update @@ -133,6 +134,9 @@ def update(self, else: request_parameters["libraries"] = [] + if encryption_key is not None: + request_parameters["encryptionKey"] = encryption_key + self._api.post("projects/update", request_parameters, data_as_json=False) def delete(self, project_id: int) -> None: diff --git a/lean/components/cloud/pull_manager.py b/lean/components/cloud/pull_manager.py index 4c88d7db..cb812458 100644 --- a/lean/components/cloud/pull_manager.py +++ b/lean/components/cloud/pull_manager.py @@ -20,11 +20,12 @@ from lean.components.util.platform_manager import PlatformManager from lean.components.util.project_manager import ProjectManager from lean.models.api import QCProject, QCLanguage, QCProjectLibrary +from lean.components.util.organization_manager import OrganizationManager from lean.models.errors import RequestFailedError from lean.models.utils import LeanLibraryReference +from lean.models.encryption import ActionType from lean.components.config.storage import safe_save - class PullManager: """The PullManager class is responsible for synchronizing cloud projects to the local drive.""" @@ -34,7 +35,8 @@ def __init__(self, project_manager: ProjectManager, project_config_manager: ProjectConfigManager, library_manager: LibraryManager, - platform_manager: PlatformManager) -> None: + platform_manager: PlatformManager, + organization_manager: OrganizationManager) -> None: """Creates a new PullManager instance. :param logger: the logger to use when printing messages @@ -50,6 +52,7 @@ def __init__(self, self._project_config_manager = project_config_manager self._library_manager = library_manager self._platform_manager = platform_manager + self._organization_manager = organization_manager self._last_file = None def _get_libraries(self, project: QCProject, @@ -90,8 +93,7 @@ def _get_libraries(self, project: QCProject, return libraries, inaccessible_libraries - def pull_projects(self, projects_to_pull: List[QCProject], - all_cloud_projects: Optional[List[QCProject]] = None) -> None: + def pull_projects(self, projects_to_pull: List[QCProject], all_cloud_projects: Optional[List[QCProject]] = None, encryption_action: Optional[ActionType]=None, encryption_key: Optional[Path]=None) -> None: """Pulls the given projects from the cloud to the local drive. This will also pull libraries referenced by the project. @@ -135,7 +137,7 @@ def pull_projects(self, projects_to_pull: List[QCProject], 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))) + projects_with_paths.append((project, self._pull_project(project, encryption_action, encryption_key))) except Exception as ex: from traceback import format_exc self._logger.debug(format_exc().strip()) @@ -147,7 +149,7 @@ def pull_projects(self, projects_to_pull: List[QCProject], self._update_local_library_references(projects_with_paths) - def _pull_project(self, project: QCProject) -> Path: + def _pull_project(self, project: QCProject, encryption_action: Optional[ActionType], encryption_key: Optional[Path]) -> Path: """Pulls a single project from the cloud to the local drive. Raises an error with a descriptive message if the project cannot be pulled. @@ -174,8 +176,17 @@ def _pull_project(self, project: QCProject) -> Path: if project_name_on_disk != project.name: self._project_manager.rename_project_and_contents(project_path_on_disk, Path.cwd() / project.name) + project_config = self._project_config_manager.get_project_config(local_project_path) + local_encryption_state = project_config.get("encrypted", False) + local_encryption_key = project_config.get("encryption-key-path", None) + if local_encryption_key is not None: + local_encryption_key = Path(local_encryption_key) + # Handle mismatch cases + 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, local_encryption_key, self._logger) + # Pull the cloud files to the local drive - self._pull_files(project, local_project_path) + self._pull_files(project, local_project_path, encryption_action, encryption_key) # Update the local project config with the latest details project_config = self._project_config_manager.get_project_config(local_project_path) @@ -185,6 +196,10 @@ def _pull_project(self, project: QCProject) -> Path: project_config.set("description", project.description) project_config.set("organization-id", project.organizationId) project_config.set("python-venv", project.leanEnvironment) + 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) @@ -193,7 +208,7 @@ def _pull_project(self, project: QCProject) -> Path: return local_project_path - def _pull_files(self, project: QCProject, local_project_path: Path) -> None: + def _pull_files(self, project: QCProject, local_project_path: Path, encryption_action: Optional[ActionType], encryption_key: Optional[Path]) -> None: """Pull the files of a single project. :param project: the cloud project of which the files need to be pulled @@ -202,7 +217,13 @@ def _pull_files(self, project: QCProject, local_project_path: Path) -> None: if not local_project_path.exists(): self._project_manager.create_new_project(local_project_path, project.language) - for cloud_file in self._api_client.files.get_all(project.projectId): + 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() + 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 if cloud_file.isLibrary: diff --git a/lean/components/cloud/push_manager.py b/lean/components/cloud/push_manager.py index 180aa767..083648c7 100644 --- a/lean/components/cloud/push_manager.py +++ b/lean/components/cloud/push_manager.py @@ -12,7 +12,7 @@ # limitations under the License. from pathlib import Path -from typing import List, Dict +from typing import List, Dict, Optional from lean.components.api.api_client import APIClient from lean.components.config.project_config_manager import ProjectConfigManager @@ -21,6 +21,7 @@ from lean.components.util.project_manager import ProjectManager from lean.models.api import QCLanguage, QCProject from lean.models.utils import LeanLibraryReference +from lean.models.encryption import ActionType class PushManager: """The PushManager class is responsible for synchronizing local projects to the cloud.""" @@ -45,7 +46,7 @@ def __init__(self, self._organization_manager = organization_manager self._cloud_projects = [] - def push_project(self, project: Path) -> None: + def push_project(self, project: Path, encryption_action: Optional[ActionType]=None, encryption_key: Optional[Path]=None) -> None: """Pushes the given project from the local drive to the cloud. It will also push every library referenced by the project and add or remove references. @@ -53,9 +54,9 @@ def push_project(self, project: Path) -> None: :param project: path to the directory containing the local project that needs to be pushed """ libraries = self._project_manager.get_project_libraries(project) - self.push_projects(libraries + [project]) + self.push_projects(libraries + [project], encryption_action, encryption_key) - def push_projects(self, projects_to_push: List[Path]) -> None: + def push_projects(self, projects_to_push: List[Path], encryption_action: Optional[ActionType]=None, encryption_key: Optional[Path]=None) -> None: """Pushes the given projects from the local drive to the cloud. It will also push every library referenced by each project and add or remove references. @@ -71,7 +72,7 @@ def push_projects(self, projects_to_push: List[Path]) -> None: relative_path = path.relative_to(Path.cwd()) try: self._logger.info(f"[{index}/{len(projects_to_push)}] Pushing '{relative_path}'") - self._push_project(path, organization_id) + self._push_project(path, organization_id, encryption_action, encryption_key) except Exception as ex: from traceback import format_exc self._logger.debug(format_exc().strip()) @@ -88,7 +89,7 @@ def _get_local_libraries_cloud_ids(self, project_dir: Path) -> List[int]: return local_libraries_cloud_ids - def _push_project(self, project_path: Path, organization_id: str, suggested_rename_path: Path = None) -> None: + def _push_project(self, project_path: Path, organization_id: str, encryption_action: Optional[ActionType], encryption_key: Optional[Path], suggested_rename_path: Path = None) -> None: """Pushes a single local project to the cloud. Raises an error with a descriptive message if the project cannot be pushed. @@ -107,6 +108,10 @@ def _push_project(self, project_path: Path, organization_id: str, suggested_rena project_config = self._project_config_manager.get_project_config(project_path) cloud_id = project_config.get("cloud-id") + local_encryption_state = project_config.get("encrypted", False) + local_encryption_key = project_config.get("encryption-key-path", None) + if local_encryption_key is not None: + local_encryption_key = Path(local_encryption_key) # check if project name is valid or if rename is required if cloud_id is not None: @@ -138,33 +143,51 @@ def _push_project(self, project_path: Path, organization_id: str, suggested_rena if cloud_project.name != project_name: # cloud project name was changed. Repeat steps to validate the new name locally. self._logger.info(f"Received new name '{cloud_project.name}' for project '{project_name}' from QuantConnect.com") - self._push_project(project_path, organization_id, Path.cwd() / cloud_project.name) + self._push_project(project_path, organization_id, encryption_action, encryption_key, Path.cwd() / cloud_project.name) return self._cloud_projects.append(cloud_project) organization_message_part = f" in organization '{organization_id}'" if organization_id is not None else "" self._logger.info(f"Successfully created cloud project '{cloud_project.name}'{organization_message_part}") + # Handle mismatch cases + 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, local_encryption_key, self._logger) + if local_encryption_state == True and encryption_key is None and local_encryption_key is not None: + encryption_key = local_encryption_key + encryption_action = ActionType.ENCRYPT if local_encryption_state else ActionType.DECRYPT # Finalize pushing by updating locally modified metadata, files and libraries - self._push_metadata(project_path, cloud_project) + self._push_metadata(project_path, cloud_project, encryption_action, encryption_key) - def _get_files(self, project: Path) -> List[Dict[str, str]]: + def _get_files(self, project: Path, encryption_action: Optional[ActionType], encryption_key: Optional[Path]) -> List[Dict[str, str]]: """Pushes the files of a local project to the cloud. :param project: the local project to push the files of """ paths = self._project_manager.get_source_files(project) - files = [ + 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() + 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(), - 'content': path.read_text(encoding="utf-8") + 'content': encrypted_content } - for path in paths + for path, encrypted_content in zip(paths, data) ] + else: + files = [ + { + 'name': path.relative_to(project).as_posix(), + 'content': path.read_text(encoding="utf-8") + } + for path in paths + ] return files - def _push_metadata(self, project: Path, cloud_project: QCProject) -> None: + def _push_metadata(self, project: Path, cloud_project: QCProject, encryption_action: Optional[ActionType], encryption_key: Optional[Path]) -> None: """Pushes local project description and parameters to the cloud. Does nothing if the cloud is already up-to-date. @@ -212,9 +235,20 @@ def _push_metadata(self, project: Path, cloud_project: QCProject) -> None: if local_lean_venv is not None and local_lean_venv != cloud_lean_venv: update_args["python_venv"] = local_lean_venv - update_args["files"] = self._get_files(project) + update_args["files"] = self._get_files(project, encryption_action, encryption_key) update_args["libraries"] = self._get_local_libraries_cloud_ids(project) + # default value + update_args["encryption_key"] = '' + if (encryption_action is not None and encryption_key is not None): + # lets check if the given key is registered with the cloud + from lean.components.util.encryption_helper import validate_encryption_key_registered_with_cloud, get_project_key_hash + validate_encryption_key_registered_with_cloud(encryption_key, self._organization_manager, self._api_client) + + if encryption_action == ActionType.ENCRYPT: + encryption_key_id = get_project_key_hash(encryption_key) + update_args["encryption_key"] = encryption_key_id + if update_args != {}: self._api_client.projects.update(cloud_project.projectId, **update_args) diff --git a/lean/components/util/encryption_helper.py b/lean/components/util/encryption_helper.py new file mode 100644 index 00000000..8a50e3b8 --- /dev/null +++ b/lean/components/util/encryption_helper.py @@ -0,0 +1,271 @@ +# QUANTCONNECT.COM - Democratizing Finance, Empowering Individuals. +# Lean CLI v1.0. Copyright 2021 QuantConnect Corporation. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import List +from pathlib import Path +from lean.components.util.logger import Logger +from base64 import b64decode, b64encode +from cryptography.hazmat.primitives import padding +from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes +from cryptography.hazmat.backends import default_backend +from cryptography.hazmat.primitives import hashes +from lean.models.encryption import ActionType +from lean.components.config.project_config_manager import ProjectConfigManager +from lean.models.api import QCProject, QCFullFile +from lean.components.config.storage import Storage +from lean.components.api.api_client import APIClient +from lean.components.util.organization_manager import OrganizationManager + +def calculate_md5(input_string: str): + """Calculate the md5 hash of a string + + :param input_string: The string to hash + :return: The md5 hash of the string + """ + from hashlib import md5 + return md5(input_string.encode()).hexdigest() + +def get_b64_encoded(key: str) -> bytes: + """Encode a string to base64 + + :param key: The string to encode + :return: The base64 encoded string + """ + return b64encode(key.encode('utf-8')) + +def get_project_key(project_key_path: Path, organization_id: str) -> str: + """Get the project key from the project key file + + :param project_key_path: The path to the project key file + :return: The project key + """ + with open(project_key_path, 'r') as f: + content = f.read() + key_for_aes = _get_fixed_length_key_from_user_full_length_key(content, organization_id.encode('utf-8')) + return key_for_aes + +def get_project_key_hash(project_key_path: Path): + """Get the MD5 hash from the project key file + + :param project_key_path: The path to the project key file + :return: The project iv + """ + with open(project_key_path, 'r', encoding='utf-8') as f: + content = f.read() + return calculate_md5(content) + +def get_project_iv(project_key_path: Path): + """Get the project iv from the project key file + + :param project_key_path: The path to the project key file + :return: The project iv + """ + key_id = get_project_key_hash(project_key_path) + return key_id[:16] + +def get_decrypted_file_content_for_local_project(project: Path, source_files: List[Path], encryption_key: Path, project_config_manager: ProjectConfigManager, organization_id: str) -> List[str]: + project_config = project_config_manager.get_project_config(project) + + # Check if the project is already encrypted + areProjectFilesAlreadyEncrypted = project_config.get('encrypted', False) + + # Check if there is mismatch of keys + _validate_key_state_for_local_project(project_config, encryption_key) + + project_key = get_project_key(encryption_key, organization_id) + project_iv = get_project_iv(encryption_key) + encoded_project_key = project_key.encode('utf-8') + encoded_project_iv = project_iv.encode('utf-8') + decrypted_data = [] + for file in source_files: + try: + # lets read and decrypt the file + with open(file, 'r') as f: + encrypted = f.read() + if not areProjectFilesAlreadyEncrypted: + decrypted = encrypted + else: + decrypted = _decrypt_file_content(encoded_project_key, encoded_project_iv, encrypted) + decrypted_data.append(decrypted) + except Exception as e: + raise RuntimeError(f"Failed to decrypt file {file} with error {e}") + return decrypted_data + +def get_encrypted_file_content_for_local_project(project: Path, source_files: List[Path], encryption_key: Path, project_config_manager: ProjectConfigManager, organization_id: str) -> List[str]: + project_config = project_config_manager.get_project_config(project) + + # Check if the project is already encrypted + areProjectFilesAlreadyEncrypted = project_config.get('encrypted', False) + + # Check if there is mismatch of keys + _validate_key_state_for_local_project(project_config, encryption_key) + + project_key = get_project_key(encryption_key, organization_id) + project_iv = get_project_iv(encryption_key) + encoded_project_key = project_key.encode('utf-8') + encoded_project_iv = project_iv.encode('utf-8') + encrypted_data: List[str] = [] + for file in source_files: + try: + # lets read and decrypt the file + with open(file, 'r') as f: + plain_text = f.read() + if areProjectFilesAlreadyEncrypted: + encrypted = plain_text + else: + encrypted = _encrypt_file_content(encoded_project_key, encoded_project_iv, plain_text.encode('utf-8')) + encrypted_data.append(encrypted) + except Exception as e: + raise RuntimeError(f"Failed to encrypt file {file} with error {e}") + return encrypted_data + +def get_and_validate_user_input_encryption_key(user_input_key: Path, project_config_encryption_key: Path) -> str: + if project_config_encryption_key is not None and Path(project_config_encryption_key).exists(): + project_config_encryption_key = Path(project_config_encryption_key) + if project_config_encryption_key is None and user_input_key is None: + raise RuntimeError("No encryption key was provided, please provide one using --key") + elif project_config_encryption_key is None: + project_config_encryption_key = user_input_key + elif user_input_key is not None and project_config_encryption_key != user_input_key: + raise RuntimeError(f"Provided encryption key ({user_input_key}) does not match the encryption key in the project ({project_config_encryption_key})") + return project_config_encryption_key + +def validate_user_inputs_for_cloud_push_pull_commands(encrypt: bool, decrypt: bool, key: Path): + if encrypt and decrypt: + raise RuntimeError(f"Cannot encrypt and decrypt at the same time.") + if key is None and (encrypt or decrypt): + raise RuntimeError(f"Encryption key is required when encrypting or decrypting.") + if key is not None and not encrypt and not decrypt: + raise RuntimeError(f"Encryption key can only be specified when encrypting or decrypting.") + +def validate_encryption_key_registered_with_cloud(user_key: Path, organization_manager: OrganizationManager, api_client: APIClient): + # lets check if the given key is registered with the cloud + organization_id = organization_manager.try_get_working_organization_id() + available_encryption_keys = api_client.encryption_keys.list(organization_id)['keys'] + encryption_key_id = get_project_key_hash(user_key) + if (not any(found_key for found_key in available_encryption_keys if found_key['hash'] == encryption_key_id)): + raise RuntimeError(f"Given encryption key is not registered with the cloud.") + +def validate_key_and_encryption_state_for_cloud_project(project: QCProject, local_project_encryption_state: bool, encryption_key: Path, local_encryption_key: Path, logger:Logger) -> None: + if not encryption_key and project.encryptionKey and local_encryption_key and local_encryption_key.exists() and get_project_key_hash(local_encryption_key) != project.encryptionKey.id: + raise RuntimeError(f"Encryption Key mismatch. Local Project Key: {local_encryption_key}. Cloud Project Key: {project.encryptionKey.name}. Please provide correct encryption key for project '{project.name}' to proceed.") + if not encryption_key and bool(project.encrypted) != bool(local_project_encryption_state): + logger.debug(f"Force Overwrite: Project encryption state mismatch. Local Project Encrypted: {bool(local_project_encryption_state)}. Cloud Project Encrypted: {bool(project.encrypted)}.") + return + if encryption_key and project.encryptionKey and get_project_key_hash(encryption_key) != project.encryptionKey.id: + raise RuntimeError(f"Encryption Key mismatch. Local Project Key hash: {get_project_key_hash(encryption_key)}. Cloud Project Key Hash: {project.encryptionKey.id}. Please provide correct encryption key for project '{project.name}' to proceed.") + +def get_appropriate_files_from_cloud_project(project: QCProject, cloud_files: List[QCFullFile], encryption_key: Path, organization_id: str, encryption_action: ActionType) -> List[QCFullFile]: + if encryption_action == ActionType.DECRYPT: + return _get_decrypted_content_from_cloud_project(project, cloud_files, encryption_key, organization_id) + return _get_encrypted_content_from_cloud_project(project, cloud_files, encryption_key, organization_id) + +def get_appropriate_files_from_local_project(project: Path, paths: List[Path], encryption_key: Path, project_config_manager: ProjectConfigManager, organization_id: str, encryption_action: ActionType) -> List[str]: + if encryption_action == ActionType.ENCRYPT: + return get_encrypted_file_content_for_local_project(project, paths, encryption_key, project_config_manager, organization_id) + return get_decrypted_file_content_for_local_project(project, paths, encryption_key,project_config_manager, organization_id) + + +def _validate_key_state_for_local_project(project_config: Storage, encryption_key: Path): + local_encryption_key_path = project_config.get('encryption-key-path', None) + if local_encryption_key_path and encryption_key and Path(local_encryption_key_path) != encryption_key: + raise RuntimeError(f"Registered encryption key {local_encryption_key_path} is different from the one provided {encryption_key}") + +def _decrypt_file_content(key: bytes, init_vector: bytes, b64_encoded_content: str) -> str: + # remove new line characters that we added during encryption + b64_encoded_content = b64_encoded_content.replace('\n', '') + b64_encoded_content = b64_encoded_content.strip() + content = b64decode(b64_encoded_content.encode('utf-8')) + # Setup module-specific classes + cipher = Cipher(algorithms.AES(key), modes.CBC(init_vector)) + decryptor = cipher.decryptor() + + plaintext = decryptor.update(content) + decryptor.finalize() + # Unpad the decrypted data + unpadder = padding.PKCS7(128).unpadder() + unpadded_data = unpadder.update(plaintext) + unpadder.finalize() + return unpadded_data.decode('utf-8').replace("\r\n", "\n") + +def _encrypt_file_content(key: bytes, init_vector: bytes, content: bytes) -> str: + plain_text = _pad(content, 16) + # Setup module-specific classes + cipher = Cipher(algorithms.AES(key), modes.CBC(init_vector)) + encryptor = cipher.encryptor() + + # Encrypt and decrypt data + cipher_text = encryptor.update(plain_text) + encryptor.finalize() + encrypted_content = b64encode(cipher_text).decode('utf-8') + # let's make it user friendly and add new lines, same as local platform + regex = r".{1,80}" + from re import findall + chunks = findall(regex, encrypted_content) + return '\n'.join(chunks) + +def _get_decrypted_content_from_cloud_project(project: QCProject, cloud_files: List[QCFullFile], encryption_key: Path, organization_id: str) -> List[QCFullFile]: + # Check if the project is already encrypted + if not project.encrypted or project.encryptionKey.id != get_project_key_hash(encryption_key): + return cloud_files + + project_key = get_project_key(encryption_key, organization_id) + project_iv = get_project_iv(encryption_key) + encoded_project_key = project_key.encode('utf-8') + encoded_project_iv = project_iv.encode('utf-8') + for cloud_file in cloud_files: + try: + decrypted = _decrypt_file_content(encoded_project_key, encoded_project_iv, cloud_file.content) + cloud_file.content = decrypted + except Exception as e: + raise RuntimeError(f"Failed to decrypt file {cloud_file} with error {e}") + return cloud_files + +def _get_encrypted_content_from_cloud_project(project: QCProject, cloud_files: List[QCFullFile], encryption_key: Path, organization_id: str) -> List[QCFullFile]: + # Check if the project is already encrypted + if project.encrypted: + if encryption_key is not None and project.encryptionKey and project.encryptionKey.id != get_project_key_hash(encryption_key): + raise RuntimeError(f"Registered encryption key {project.encryptionKey.id} is different from the one provided {encryption_key}") + return cloud_files + + project_key = get_project_key(encryption_key, organization_id) + project_iv = get_project_iv(encryption_key) + encoded_project_key = project_key.encode('utf-8') + encoded_project_iv = project_iv.encode('utf-8') + for cloud_file in cloud_files: + try: + encrypted = _encrypt_file_content(encoded_project_key, encoded_project_iv, cloud_file.content.encode('utf-8')) + cloud_file.content = encrypted + except Exception as e: + raise RuntimeError(f"Failed to decrypt file {cloud_file} with error {e}") + return cloud_files + +def _pad(data, block_size): + """Padding function for AES encryption + + :param data: data to pad + :param block_size: required block size + :return: padded data + """ + padding_length = block_size - (len(data) % block_size) + return data + bytes([padding_length] * padding_length) + +def _get_fixed_length_key_from_user_full_length_key(password: str, salt: bytes): + from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC + kdf = PBKDF2HMAC( + algorithm=hashes.SHA256(), + length=16, + salt=salt, + iterations=100000, + backend=default_backend() + ) + return kdf.derive(password.encode()).hex() + + diff --git a/lean/container.py b/lean/container.py index 226f8ebb..3f82d5f6 100644 --- a/lean/container.py +++ b/lean/container.py @@ -120,7 +120,8 @@ def initialize(self, self.project_manager, self.project_config_manager, self.library_manager, - self.platform_manager) + self.platform_manager, + self.organization_manager) self.push_manager = push_manager if not push_manager: diff --git a/lean/models/api.py b/lean/models/api.py index ea4916e3..b77f741a 100644 --- a/lean/models/api.py +++ b/lean/models/api.py @@ -23,6 +23,10 @@ # The keys of properties are not changed, so they don't obey the rest of the project's naming conventions +class ProjectEncryptionKey(WrappedBaseModel): + id: str + name: str + class QCCollaborator(WrappedBaseModel): id: int uid: int @@ -79,6 +83,8 @@ class QCProject(WrappedBaseModel): leanEnvironment: int parameters: List[QCParameter] libraries: List[QCProjectLibrary] + encrypted: Optional[bool] = False + encryptionKey: Optional[ProjectEncryptionKey] = None @validator("parameters", pre=True) def process_parameters_dict(cls, value: Any) -> Any: diff --git a/lean/models/encryption.py b/lean/models/encryption.py new file mode 100644 index 00000000..e54f372a --- /dev/null +++ b/lean/models/encryption.py @@ -0,0 +1,18 @@ +# QUANTCONNECT.COM - Democratizing Finance, Empowering Individuals. +# Lean CLI v1.0. Copyright 2021 QuantConnect Corporation. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from enum import Enum + +class ActionType(Enum): + ENCRYPT = "encrypt" + DECRYPT = "decrypt" diff --git a/setup.py b/setup.py index fa758d99..84aaa974 100644 --- a/setup.py +++ b/setup.py @@ -58,7 +58,8 @@ def get_stubs_version_range() -> str: "joblib>=1.1.0", "wrapt~=1.14.1", "setuptools", - f"quantconnect-stubs{get_stubs_version_range()}" + f"quantconnect-stubs{get_stubs_version_range()}", + "cryptography~=41.0.4" ] setup( diff --git a/tests/commands/cloud/test_pull.py b/tests/commands/cloud/test_pull.py index 5dcba744..d561f4e9 100644 --- a/tests/commands/cloud/test_pull.py +++ b/tests/commands/cloud/test_pull.py @@ -12,13 +12,25 @@ # limitations under the License. from unittest import mock - +from pathlib import Path +from datetime import datetime from click.testing import CliRunner - +from lean.models.api import QCFullFile from lean.commands import lean from lean.container import container +from tests.conftest import initialize_container from tests.test_helpers import create_api_project, create_fake_lean_cli_directory +from lean.components.util.encryption_helper import get_project_key_hash + +def init_container(**kwargs) -> None: + organization_manager = mock.Mock() + organization_manager.get_working_organization_id = mock.MagicMock(return_value="abc") + organization_manager.try_get_working_organization_id = mock.MagicMock(return_value="abc") + + if "organization_manager_to_use" not in kwargs: + kwargs["organization_manager_to_use"] = organization_manager + initialize_container(**kwargs) def test_cloud_pull_pulls_all_non_bootcamp_projects_when_no_options_given() -> None: create_fake_lean_cli_directory() @@ -40,7 +52,7 @@ def test_cloud_pull_pulls_all_non_bootcamp_projects_when_no_options_given() -> N assert result.exit_code == 0 - pull_manager.pull_projects.assert_called_once_with(cloud_projects[:3], cloud_projects) + pull_manager.pull_projects.assert_called_once_with(cloud_projects[:3], cloud_projects, None, None) def test_cloud_pull_pulls_all_projects_when_pull_bootcamp_option_given() -> None: @@ -63,7 +75,7 @@ def test_cloud_pull_pulls_all_projects_when_pull_bootcamp_option_given() -> None assert result.exit_code == 0 - pull_manager.pull_projects.assert_called_once_with(cloud_projects, cloud_projects) + pull_manager.pull_projects.assert_called_once_with(cloud_projects, cloud_projects, None, None) def test_cloud_pull_pulls_project_by_id() -> None: @@ -87,7 +99,7 @@ def test_cloud_pull_pulls_project_by_id() -> None: assert result.exit_code == 0 - pull_manager.pull_projects.assert_called_once_with([project_to_pull], None) + pull_manager.pull_projects.assert_called_once_with([project_to_pull], None, None, None) def test_cloud_pull_pulls_project_by_name() -> None: @@ -110,7 +122,7 @@ def test_cloud_pull_pulls_project_by_name() -> None: assert result.exit_code == 0 - pull_manager.pull_projects.assert_called_once_with([cloud_projects[0]], cloud_projects) + pull_manager.pull_projects.assert_called_once_with([cloud_projects[0]], cloud_projects, None, None) def test_cloud_pull_aborts_when_project_input_matches_no_cloud_projects() -> None: @@ -134,3 +146,363 @@ def test_cloud_pull_aborts_when_project_input_matches_no_cloud_projects() -> Non assert result.exit_code != 0 pull_manager.pull_projects.assert_not_called() + +def test_cloud_pull_aborts_when_encrypting_without_key_given() -> None: + create_fake_lean_cli_directory() + + pull_manager = mock.Mock() + container.pull_manager = pull_manager + + (Path.cwd() / "Empty Project").mkdir() + + result = CliRunner().invoke(lean, ["cloud", "pull", "--project", "Empty Project", "--encrypt"]) + + assert result.exit_code != 0 + + pull_manager.pull_projects.assert_not_called() + +def test_cloud_pull_aborts_when_decrypting_without_key_given() -> None: + create_fake_lean_cli_directory() + + pull_manager = mock.Mock() + container.pull_manager = pull_manager + + (Path.cwd() / "Empty Project").mkdir() + + result = CliRunner().invoke(lean, ["cloud", "pull", "--project", "Empty Project", "--decrypt"]) + + assert result.exit_code != 0 + + pull_manager.pull_projects.assert_not_called() + + +def test_cloud_pull_receives_encrypted_files_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_projects = [create_api_project(1, "Python Project")] + + api_client = mock.Mock() + api_client.projects.get_all.return_value = cloud_projects + + initial_source_files = container.project_manager.get_source_files(project_path) + fake_cloud_files = [QCFullFile(name=file.name, content=file.read_text(), modified=datetime.now(), isLibrary=False) + for file in initial_source_files] + 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, "--encrypt", "--key", encryption_file_path]) + + assert result.exit_code == 0 + + source_files = container.project_manager.get_source_files(project_path) + expected_encrypted_files = _get_expected_encrypted_files_content() + for file in source_files: + assert expected_encrypted_files[file.name].strip() == file.read_text().strip() + +def test_cloud_pull_aborts_receiving_encrypted_files_when_cloud_file_encrypted_with_key_x_and_given_key_y() -> None: + create_fake_lean_cli_directory() + + project_path = Path.cwd() / "Python Project" + + encryption_file_path_x = project_path / "encryption_x.txt" + encryption_file_path_x.write_text("KtSwJtq5a4uuQmxbPqcCP3d8yMRz5TZxDBAKy7kGwPcvcvsNBdCprGYwSBN8ntJa5JNNYHTB2GrBpAbkA38kCdnceegffZH7") + + encryption_file_path_y = project_path / "encryption_y.txt" + encryption_file_path_y.write_text("Jtq5a4uuQmxbPqcCP3d8yMRz5TZxDBAKy7kGwPcvcvsNBdCprGYwSBN8ntJa5JNNYHTB2GrBpAbkA38kCdnceegffZH7") + key_hash_y = get_project_key_hash(encryption_file_path_y) + + api_client = mock.Mock() + cloud_project = create_api_project(1, "Python Project", encrypted=True, encryptionKey={"name":"test", "id": key_hash_y}) + cloud_projects = [cloud_project] + api_client.projects.get_all.return_value = cloud_projects + api_client.projects.get.return_value = cloud_project + + initial_source_files = container.project_manager.get_source_files(project_path) + file_contents_map = {file.name: file.read_text() for file in initial_source_files} + fake_cloud_files = [QCFullFile(name=file.name, content=file.read_text(), modified=datetime.now(), isLibrary=False) + for file in initial_source_files] + 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_x)) + project_config.set("cloud-id", 1) + + result = CliRunner().invoke(lean, ["cloud", "pull", "--project", project_path.name, "--encrypt", "--key", encryption_file_path_x]) + + assert result.exit_code == 0 + source_files = container.project_manager.get_source_files(project_path) + for file in source_files: + assert file_contents_map[file.name].strip() == file.read_text().strip() + +def test_cloud_pull_turns_on_encryption_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") + # Keys API Data + 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}) + cloud_projects = [cloud_project] + + api_client = mock.Mock() + api_client.projects.get_all.return_value = cloud_projects + api_client.projects.get.return_value = cloud_project + + initial_source_files = container.project_manager.get_source_files(project_path) + fake_cloud_files = [QCFullFile(name=file.name, content=file.read_text(), modified=datetime.now(), isLibrary=False) + for file in initial_source_files] + 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("cloud-id", 1) + assert project_config.get("encrypted", False) == False + 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 + +def test_cloud_pull_receives_decrypted_files_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") + # Keys API Data + 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}) + cloud_projects = [cloud_project] + + api_client = mock.Mock() + api_client.projects.get_all.return_value = cloud_projects + api_client.projects.get.return_value = cloud_project + + encrypted_source_files = _get_expected_encrypted_files_content() + initial_source_files = container.project_manager.get_source_files(project_path) + file_contents_map = {file.name: file.read_text() for file in initial_source_files} + # replace the content of the files with the encrypted content and verify later that they are decrypted. + for file in initial_source_files: + file.write_text(encrypted_source_files[file.name]) + fake_cloud_files = [QCFullFile(name=name, content=content, modified=datetime.now(), isLibrary=False) + for name, content in encrypted_source_files.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("cloud-id", 1) + result = CliRunner().invoke(lean, ["cloud", "pull", "--project", project_path.name, "--decrypt", "--key", encryption_file_path]) + + assert result.exit_code == 0 + + source_files = container.project_manager.get_source_files(project_path) + for file in source_files: + assert file_contents_map[file.name].strip() == file.read_text().strip() + +def test_cloud_pull_turns_off_encryption_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") + # Keys API Data + 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}) + cloud_projects = [cloud_project] + + api_client = mock.Mock() + api_client.projects.get_all.return_value = cloud_projects + api_client.projects.get.return_value = cloud_project + + encrypted_source_files = _get_expected_encrypted_files_content() + fake_cloud_files = [QCFullFile(name=name, content=content, modified=datetime.now(), isLibrary=False) + for name, content in encrypted_source_files.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 + +def test_cloud_pull_decrypted_files_when_local_files_in_encrypted_state_and_cloud_project_in_decrypted_state_without_key_given() -> None: + create_fake_lean_cli_directory() + + project_path = Path.cwd() / "Python Project" + + encryption_file_path_x = project_path / "encryption_x.txt" + encryption_file_path_x.write_text("KtSwJtq5a4uuQmxbPqcCP3d8yMRz5TZxDBAKy7kGwPcvcvsNBdCprGYwSBN8ntJa5JNNYHTB2GrBpAbkA38kCdnceegffZH7") + + api_client = mock.Mock() + cloud_projects = [create_api_project(1, "Python Project")] + api_client.projects.get_all.return_value = cloud_projects + + initial_source_files = container.project_manager.get_source_files(project_path) + file_contents_map = {file.name: file.read_text() for file in initial_source_files} + fake_cloud_files = [QCFullFile(name=file.name, content=file.read_text(), modified=datetime.now(), isLibrary=False) + for file in initial_source_files] + api_client.files.get_all = mock.MagicMock(return_value=fake_cloud_files) + encrypted_source_files = _get_expected_encrypted_files_content() + # replace the content of the files with the encrypted content and verify later that they are decrypted. + for file in initial_source_files: + file.write_text(encrypted_source_files[file.name]) + + 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_x)) + project_config.set("cloud-id", 1) + + result = CliRunner().invoke(lean, ["cloud", "pull", "--project", project_path.name]) + + assert result.exit_code == 0 + source_files = container.project_manager.get_source_files(project_path) + for file in source_files: + assert file_contents_map[file.name].strip() == file.read_text().strip() + project_config = container.project_config_manager.get_project_config(project_path) + assert project_config.get("encrypted", False) == False + +def test_cloud_pull_encrypts_when_local_files_in_decrypted_state_and_cloud_project_in_encrypted_state_without_key_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") + # Keys API Data + 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}) + cloud_projects = [cloud_project] + + api_client = mock.Mock() + api_client.projects.get_all.return_value = cloud_projects + api_client.projects.get.return_value = cloud_project + + encrypted_source_files = _get_expected_encrypted_files_content() + fake_cloud_files = [QCFullFile(name=name, content=content, modified=datetime.now(), isLibrary=False) + for name, content in encrypted_source_files.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("cloud-id", 1) + result = CliRunner().invoke(lean, ["cloud", "pull", "--project", project_path.name]) + + assert result.exit_code == 0 + + source_files = container.project_manager.get_source_files(project_path) + expected_encrypted_files = _get_expected_encrypted_files_content() + for file in source_files: + assert expected_encrypted_files[file.name].strip() == file.read_text().strip() + +def test_cloud_pull_aborts_when_local_files_in_encrypted_state_with_key_x_and_cloud_project_in_encrypted_state_with_key_y() -> None: + create_fake_lean_cli_directory() + + project_path = Path.cwd() / "Python Project" + + encryption_file_path_x = project_path / "encryption_x.txt" + encryption_file_path_x.write_text("KtSwJtq5a4uuQmxbPqcCP3d8yMRz5TZxDBAKy7kGwPcvcvsNBdCprGYwSBN8ntJa5JNNYHTB2GrBpAbkA38kCdnceegffZH7") + + encryption_file_path_y = project_path / "encryption_y.txt" + encryption_file_path_y.write_text("Jtq5a4uuQmxbPqcCP3d8yMRz5TZxDBAKy7kGwPcvcvsNBdCprGYwSBN8ntJa5JNNYHTB2GrBpAbkA38kCdnceegffZH7") + + # Keys API Data + key_hash_y = get_project_key_hash(encryption_file_path_y) + + api_client = mock.Mock() + cloud_project = create_api_project(1, "Python Project", encrypted=True, encryptionKey={"name":"test", "id": key_hash_y}) + cloud_projects = [cloud_project] + api_client.projects.get_all.return_value = cloud_projects + api_client.projects.get.return_value = cloud_project + + initial_source_files = container.project_manager.get_source_files(project_path) + file_contents_map = {file.name: file.read_text() for file in initial_source_files} + fake_cloud_files = [QCFullFile(name=file.name, content=file.read_text(), modified=datetime.now(), isLibrary=False) + for file in initial_source_files] + 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_x)) + project_config.set("cloud-id", 1) + result = CliRunner().invoke(lean, ["cloud", "pull", "--project", project_path.name]) + + assert result.exit_code == 0 + source_files = container.project_manager.get_source_files(project_path) + for file in source_files: + assert file_contents_map[file.name].strip() == file.read_text().strip() + + +def _get_expected_encrypted_files_content() -> dict: + return { + "main.py": + """UpMdqgoXS1tgqGgy6nKkmlxrOV7ikoc5oJAmS+pcMcmD0qsfJq5GE/yvdg9mucXrhfgLjD7of3YalHYC +mJcVeO1VSlHCA3oNq5kS82YV4Rt0KL0IApPXlV7yAvJW/SbqcziwPXb/OSmHroQfaD9IJWJdNXrcjfs1 +CAXq8mV0BX0OipDTnQL2EfHCDNn3PL7lJPg0lYOKjSuMf6QgvFWElT6Kw0UDiLtKlU6jJwsKxeJDDRAw +KkMJVrrIsxk5g92ERCquoTmYHVm1lf0xdsM+vTqSRxwucEfaoi7DgA9SkJKgaVoDgx90ScqssPzdlsgJ +zIUdPbHO1GN1D5ZfHXU+7Sww4pHDh6QTyLYnuOHnPQ8IuXYDWQPdw9QaHrbeuWqqwyFl859Tra6RZFlD +136K2CZ6PcClbGeeGO2NmGPRFuwHH5B5aSfUs8Y3OMdYnZylXtuWAMPU6xEZIWfGHIOZlQZ7s06rvfYT +CfMHdkFCRi9QnI1vRbqO4yhBDB/BbAG9W7wt/FX1f+1hHk3xpaP+MLufoN48xXSxudIV975Jxdoic8YK +P/6UZ5FWDVTZDRzJ7e9hiS5LOr410G8PqDOr0JuNap3t2FvGLuuR1CFfXsHm2QcsaJgiODzAa/xSW2vq +ffQPkZxz3bhGpuJXoWe18B92m0/Scih5J224QMXoZJYs6UYEl6eQX67mlGSSntotykByWB7upnPFDPSU +l9if+gA+GrW3hyVhDSY/5W9mUA6mw7lelE0b+0cAfuEczDXInQeMbQjAtEOJufIQl7NSZVzc4BWOFoXS +VK3OnsQnbg2f9WN7mQ4UE2xVoVAOudDPSXh21HjhUa9DbqrFy4FQN66G8LgbF6Sw58d+HfCIPOVoNl1f +mMJBPythF0UjBYwle+2EpDC/9WdsVrm02OCEehbs4gfG0vOTtmwMY0Y4Urpxs01ATOfoZEg5gRPMFRD7 +zFhH4MGKrPdBA3NXnvRozJsdHRBEVPhmuhtideHMKs68o8BwDLJ8oY6dyyxpjlFhGlHELGDTo7ljCRtR +6d5M+DpJn9vEZ9xWFe60vkgIq3dpn3SwPnrrsd2Ee5TwagDF2Bj8EhMVhshlxjrNR9NNguYgZVFUAJEM +InB4pibT1Q0FLnwFOp84t8offZ5w+hmGm7Jggs99qmx2Tmu4u4067A==""" +, +"research.ipynb": + """NIiAgzU8gzaJ3YEIWysBh0e0xxm9rpAWDE4Pir/wzKtla/wcbs98GU5cdOxgd7Gjlcu0zNbFzE5x8Eyx +qSuh3cQU17xQSisPxvjfDD/h2z9AnFT1jD1Vhc+Nn5ngwpgCA6P5fHT4VhPgaKDp7r9zc8pAURcSd04M +2dmHtGs89it3QFrYNgYAvN4kIYVZuROhnUSiYN+y7kzWzLKLIK3a+y6R4ibr9ju5S0DCtIS87MwrHFbi +NmS0mqnlYmBqLPVijrgAJYu4YmqrOb+LxwQSUXL80UsgddUFtgKWDKGyTVFs0o92/x9OT2XTCcMOhhQg +X/6h5c99rP5W3GqkcDaaiueShSD43u3LXnijK6yugUqALcMIBLxF4Iczq3xLov9MfrdgDUaxFEOu0dDK +JpecJeSx9TxVUxzcWvn3alUycScBRV1w8VH0Wa4Lf9vc1YIOe9ITy/hLW2+QEBAHxrhYd7vq1MVVDdRI +lbXtLVMOVI9YYBNEunf3drQm6fCnKv+arnt/rYSIQDvhrlqem6qKS3JE4V8gVzXQDefTewx3/hppDOxL ++ccikaUcVkIAULF1Jwp9qxCpCsiz5vLXynD47pf3mhS2FC6Dd1g30xXUVNeTmpRE7TfS/gHzkyDrtTvC +LBvBnkImsJQjDCJl/NzLlFRh4wiY4SL/bTdxL2YZJebJ1zdw3PXQhvWvgztG0wmIRE/U+1xv95gc5w5x +7N2WT7F9KVHI6NrtrcU6c5hWo4q/QEBxb5ETSyNpgsVHvKxKblPGslki/aH6WOexu7tSzbuA9rSwVgiS +NWh9y3qZ6aMOMq3dDAr4wGBkkGQsirasCEt3YEa9rG5CDH8JYApCeZk39zt+oqjQ4DcYprgW/q87fFRB +thxY4cM3BGG/IIpX9oAxvFNe7PSXIx3BBkyhqiJqJRhS3YfvasBWFnJd11h9JXVsKJudVBoJvwrzxldU +Td2mSr1d//xkIe6TeqPK3ivofMoCNoDYyO57juki8ban9Z8q0lpgShWBM8zUPNnezctwJZIp6aVaVdfC +nZtG49f06w1IRC0dOj6yM+adiEPx/vKxrCn1sqd0eoAeKIs5iEm398YJlt9G0eubSSpoDC6iH1GjwH8A +NFNRoKncB4c3ocmajgNda4o32UfpwvNr+fKLmpP+57dulpA09i/FqwF2L8dDNm4JSTjfNdeooM9TIDsG +b3CIItqmNnve8fryEykjE8WgJaP70rnsXRY6WHLivy7OTNMKG8Ij9MOyx44Wq3TWY/omRDjYgfCO4QRZ +vTtlnIMp8VtLlBgxRc6TBPZutYZXDJyY4N4n2Jpspl3KUbe3iHNm6nwl0SCwrIl3aYbCs55BP/LOGXLx +dNIZzkxJhnfLtomrxO4QgcUnjNEWYJ1GvRAplTa99RFFYXOuMQID4+UH+6RM2e6cqBafv+2+lpetfbBM +gMOA+qNdar9Y+b+vLuSz2ak0MCn3h8cAAlXgKNpo4shbgn+R4TViLGfhUZbcSG9VZ0UEBYyg5nxXggqh +fPkc80yGC7lwo0XmrQrp62IaQy3iB2MYWB29KCX9MdghcobfhmW3BEtYqDRLE61/MWYi9scowMkf0TxT +DyMRNS9qVwJ+2gS44xD+lQmQuE1aM+rBmRkVKWA+fB5HzZe6yyN8GkM4D23hcw+DBs9opXvq6QyGd8Le +KqLYkbp7Be44+d6mie1aGp0hRUjF9LzNapPdnYusx4QwoIlt0wYTr1waqAYMG73PhDI36B6oa4tOwuSj +86Wje1jEfUK01LVuVYloclnG3fTQxQlUcrRyQNCbrg0M5elYWbjYybx51/gb//9McGYXW0SWltNlvLvL +SUroK4+Yhqfk5MVeascvM94fFBfieT3erU8D+iCiaVvZGcuE5mE0bLBquXpMuO4fQ06XVI/NzlhVILrV +yXa/cl9WaYDFPqM+eFVJfg==""" +} diff --git a/tests/commands/cloud/test_push.py b/tests/commands/cloud/test_push.py index f66c3813..a06ae087 100644 --- a/tests/commands/cloud/test_push.py +++ b/tests/commands/cloud/test_push.py @@ -22,11 +22,13 @@ from lean.models.api import QCFullFile from tests.test_helpers import create_fake_lean_cli_directory, create_api_project from tests.conftest import initialize_container - +from lean.container import container +from lean.components.util.encryption_helper import get_project_key_hash def init_container(**kwargs) -> None: organization_manager = mock.Mock() organization_manager.get_working_organization_id = mock.MagicMock(return_value="abc") + organization_manager.try_get_working_organization_id = mock.MagicMock(return_value="abc") if "organization_manager_to_use" not in kwargs: kwargs["organization_manager_to_use"] = organization_manager @@ -80,7 +82,7 @@ def test_cloud_push_pushes_single_project_when_project_option_given() -> None: assert result.exit_code == 0 - push_manager.push_project.assert_called_once_with(Path.cwd() / "Python Project") + push_manager.push_project.assert_called_once_with(Path.cwd() / "Python Project", None, None) def test_cloud_push_aborts_when_given_directory_is_not_lean_project() -> None: @@ -114,7 +116,7 @@ def test_cloud_push_aborts_when_given_directory_does_not_exist() -> None: def test_cloud_push_updates_lean_config() -> None: create_fake_lean_cli_directory() - + project_path = Path.cwd() / "Python Project" cloud_project = create_api_project(1, "Python Project") api_client = mock.Mock() api_client.projects.create = mock.MagicMock(return_value=cloud_project) @@ -125,22 +127,346 @@ def test_cloud_push_updates_lean_config() -> None: api_client.projects.get_all = mock.MagicMock(return_value=[cloud_project]) api_client.projects.get = mock.MagicMock(return_value=create_api_project(1, "Python Project")) - project_config = mock.Mock() - project_config.get = mock.MagicMock(side_effect=[None, "Python", "", {}, -1, None, []]) + init_container(api_client_to_use=api_client) - project_config_manager = mock.Mock() - project_config_manager.get_project_config = mock.MagicMock(return_value=project_config) + project_config = container.project_config_manager.get_project_config(project_path) + assert project_config.get("organization-id", None) == None + assert cloud_project.organizationId == "123" - project_manager = mock.Mock() - project_manager.get_source_files = mock.MagicMock(return_value=[]) - project_manager.get_project_libraries = mock.MagicMock(return_value=[]) + result = CliRunner().invoke(lean, ["cloud", "push", "--project", "Python Project"]) - push_manager = PushManager(mock.Mock(), api_client, project_manager, project_config_manager, mock.Mock()) + assert result.exit_code == 0 - init_container(push_manager_to_use=push_manager, api_client_to_use=api_client) + project_config = container.project_config_manager.get_project_config(project_path) + assert project_config.get("organization-id", None) == "123" - result = CliRunner().invoke(lean, ["cloud", "push", "--project", "Python Project"]) + +def test_cloud_push_aborts_when_encrypting_without_key_given() -> None: + create_fake_lean_cli_directory() + + push_manager = mock.Mock() + init_container(push_manager_to_use=push_manager) + + (Path.cwd() / "Empty Project").mkdir() + + result = CliRunner().invoke(lean, ["cloud", "push", "--project", "Empty Project", "--encrypt"]) + + assert result.exit_code != 0 + + push_manager.push_projects.assert_not_called() + +def test_cloud_push_aborts_when_decrypting_without_key_given() -> None: + create_fake_lean_cli_directory() + + push_manager = mock.Mock() + init_container(push_manager_to_use=push_manager) + + (Path.cwd() / "Empty Project").mkdir() + + result = CliRunner().invoke(lean, ["cloud", "push", "--project", "Empty Project", "--decrypt"]) + + assert result.exit_code != 0 + + push_manager.push_projects.assert_not_called() + + +def test_cloud_push_sends_encrypted_files_and_turns_on_encryption_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") + # Keys API Data + key_hash = get_project_key_hash(encryption_file_path) + keys_api_data = {'keys': [{'name': 'test', 'hash': key_hash}]} + + api_client = mock.Mock() + api_client.encryption_keys.list = mock.MagicMock(return_value=keys_api_data) + cloud_project = create_api_project(1, "Python Project") + api_client.projects.create = mock.MagicMock(return_value=cloud_project) + fake_cloud_files = [QCFullFile(name="file.py", content="testing", modified=datetime.now(), isLibrary=False)] + 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", False) + + result = CliRunner().invoke(lean, ["cloud", "push", "--project", project_path, "--encrypt", "--key", encryption_file_path]) + + assert result.exit_code == 0 + expected_encrypted_files = _get_expected_encrypted_files_content() + # verify that the 'encryption_key' is being set to turn on the encryption in the cloud. + expected_arguments = { + "name": "Python Project", + "description": "", + "files":[{'name': name, 'content': content} for name, content in expected_encrypted_files.items()], + "libraries": [], + "encryption_key": key_hash + } + api_client.projects.update.assert_called_once_with(1, **expected_arguments) + +def test_cloud_push_aborts_sending_encrypted_files_when_local_file_encrypted_with_key_x_and_given_key_y() -> None: + create_fake_lean_cli_directory() + + project_path = Path.cwd() / "Python Project" + + encryption_file_path_x = project_path / "encryption_x.txt" + encryption_file_path_x.write_text("KtSwJtq5a4uuQmxbPqcCP3d8yMRz5TZxDBAKy7kGwPcvcvsNBdCprGYwSBN8ntJa5JNNYHTB2GrBpAbkA38kCdnceegffZH7") + + encryption_file_path_y = project_path / "encryption_y.txt" + encryption_file_path_y.write_text("Jtq5a4uuQmxbPqcCP3d8yMRz5TZxDBAKy7kGwPcvcvsNBdCprGYwSBN8ntJa5JNNYHTB2GrBpAbkA38kCdnceegffZH7") + + # Keys API Data + key_hash = get_project_key_hash(encryption_file_path_y) + keys_api_data = {'keys': [{'name': 'test', 'hash': key_hash}]} + + api_client = mock.Mock() + api_client.encryption_keys.list = mock.MagicMock(return_value=keys_api_data) + cloud_project = create_api_project(1, "Python Project") + api_client.projects.create = mock.MagicMock(return_value=cloud_project) + fake_cloud_files = [QCFullFile(name="file.py", content="testing", modified=datetime.now(), isLibrary=False)] + 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_x)) + + result = CliRunner().invoke(lean, ["cloud", "push", "--project", project_path, "--encrypt", "--key", encryption_file_path_y]) + + assert result.exit_code == 0 + api_client.projects.update.assert_not_called() + +def test_cloud_push_sends_decrypted_files_and_turns_off_encryption_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") + # Keys API Data + key_hash = get_project_key_hash(encryption_file_path) + keys_api_data = {'keys': [{'name': 'test', 'hash': key_hash}]} + + api_client = mock.Mock() + api_client.encryption_keys.list = mock.MagicMock(return_value=keys_api_data) + cloud_project = create_api_project(1, "Python Project") + api_client.projects.create = mock.MagicMock(return_value=cloud_project) + fake_cloud_files = [QCFullFile(name="file.py", content="testing", modified=datetime.now(), isLibrary=False)] + 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", False) + + result = CliRunner().invoke(lean, ["cloud", "push", "--project", project_path, "--decrypt", "--key", encryption_file_path]) + + assert result.exit_code == 0 + # verify that the encryption key is set to empty string to turn off the encryption. + expected_arguments = { + "name": "Python Project", + "description": "", + "files": [{'name': 'main.py', 'content': '# region imports\nfrom AlgorithmImports import *\n# endregion\n\nclass PythonProject(QCAlgorithm):\n\n def Initialize(self):\n # Locally Lean installs free sample data, to download more data please visit https://www.quantconnect.com/docs/v2/lean-cli/datasets/downloading-data\n self.SetStartDate(2013, 10, 7) # Set Start Date\n self.SetEndDate(2013, 10, 11) # Set End Date\n self.SetCash(100000) # Set Strategy Cash\n self.AddEquity("SPY", Resolution.Minute)\n\n def OnData(self, data: Slice):\n """OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.\n Arguments:\n data: Slice object keyed by symbol containing the stock data\n """\n if not self.Portfolio.Invested:\n self.SetHoldings("SPY", 1)\n self.Debug("Purchased Stock")\n'}, {'name': 'research.ipynb', 'content': '{\n "cells": [\n {\n "cell_type": "markdown",\n "metadata": {},\n "source": [\n "![QuantConnect Logo](https://cdn.quantconnect.com/web/i/icon.png)\\n",\n "
"\n ]\n },\n {\n "cell_type": "code",\n "execution_count": null,\n "metadata": {},\n "outputs": [],\n "source": [\n "# QuantBook Analysis Tool \\n",\n "# For more information see [https://www.quantconnect.com/docs/v2/our-platform/research/getting-started]\\n",\n "qb = QuantBook()\\n",\n "spy = qb.AddEquity(\\"SPY\\")\\n",\n "# Locally Lean installs free sample data, to download more data please visit https://www.quantconnect.com/docs/v2/lean-cli/datasets/downloading-data \\n",\n "qb.SetStartDate(2013, 10, 11)\\n",\n "history = qb.History(qb.Securities.Keys, 360, Resolution.Daily)\\n",\n "\\n",\n "# Indicator Analysis\\n",\n "bbdf = qb.Indicator(BollingerBands(30, 2), spy.Symbol, 360, Resolution.Daily)\\n",\n "bbdf.drop(\'standarddeviation\', axis=1).plot()"\n ]\n }\n ],\n "metadata": {\n "kernelspec": {\n "display_name": "Python 3",\n "language": "python",\n "name": "python3"\n }\n },\n "nbformat": 4,\n "nbformat_minor": 2\n}\n'}], + "libraries": [], + "encryption_key": '' + } + api_client.projects.update.assert_called_once_with(1, **expected_arguments) + +def test_cloud_push_sends_decrypted_files_when_project_in_encrypted_state_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") + # Keys API Data + key_hash = get_project_key_hash(encryption_file_path) + keys_api_data = {'keys': [{'name': 'test', 'hash': key_hash}]} + + api_client = mock.Mock() + api_client.encryption_keys.list = mock.MagicMock(return_value=keys_api_data) + cloud_project = create_api_project(1, "Python Project") + api_client.projects.create = mock.MagicMock(return_value=cloud_project) + fake_cloud_files = [QCFullFile(name="file.py", content="testing", modified=datetime.now(), isLibrary=False)] + 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)) + # update the project with encrypted files so that it's in encrypted state + expected_encrypted_files = _get_expected_encrypted_files_content() + source_files = container.project_manager.get_source_files(project_path) + for source_file in source_files: + source_file.write_text(expected_encrypted_files[source_file.name]) + + result = CliRunner().invoke(lean, ["cloud", "push", "--project", project_path, "--decrypt", "--key", encryption_file_path]) + + assert result.exit_code == 0 + # verify that the encryption key is set to empty string to turn off the encryption. + expected_arguments = { + "name": "Python Project", + "description": "", + "files": [{'name': 'main.py', 'content': '# region imports\nfrom AlgorithmImports import *\n# endregion\n\nclass PythonProject(QCAlgorithm):\n\n def Initialize(self):\n # Locally Lean installs free sample data, to download more data please visit https://www.quantconnect.com/docs/v2/lean-cli/datasets/downloading-data\n self.SetStartDate(2013, 10, 7) # Set Start Date\n self.SetEndDate(2013, 10, 11) # Set End Date\n self.SetCash(100000) # Set Strategy Cash\n self.AddEquity("SPY", Resolution.Minute)\n\n def OnData(self, data: Slice):\n """OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.\n Arguments:\n data: Slice object keyed by symbol containing the stock data\n """\n if not self.Portfolio.Invested:\n self.SetHoldings("SPY", 1)\n self.Debug("Purchased Stock")\n'}, {'name': 'research.ipynb', 'content': '{\n "cells": [\n {\n "cell_type": "markdown",\n "metadata": {},\n "source": [\n "![QuantConnect Logo](https://cdn.quantconnect.com/web/i/icon.png)\\n",\n "
"\n ]\n },\n {\n "cell_type": "code",\n "execution_count": null,\n "metadata": {},\n "outputs": [],\n "source": [\n "# QuantBook Analysis Tool \\n",\n "# For more information see [https://www.quantconnect.com/docs/v2/our-platform/research/getting-started]\\n",\n "qb = QuantBook()\\n",\n "spy = qb.AddEquity(\\"SPY\\")\\n",\n "# Locally Lean installs free sample data, to download more data please visit https://www.quantconnect.com/docs/v2/lean-cli/datasets/downloading-data \\n",\n "qb.SetStartDate(2013, 10, 11)\\n",\n "history = qb.History(qb.Securities.Keys, 360, Resolution.Daily)\\n",\n "\\n",\n "# Indicator Analysis\\n",\n "bbdf = qb.Indicator(BollingerBands(30, 2), spy.Symbol, 360, Resolution.Daily)\\n",\n "bbdf.drop(\'standarddeviation\', axis=1).plot()"\n ]\n }\n ],\n "metadata": {\n "kernelspec": {\n "display_name": "Python 3",\n "language": "python",\n "name": "python3"\n }\n },\n "nbformat": 4,\n "nbformat_minor": 2\n}\n'}], + "libraries": [], + "encryption_key": '' + } + api_client.projects.update.assert_called_once_with(1, **expected_arguments) + +def test_cloud_push_encrypts_when_local_files_in_encrypted_state_and_cloud_project_in_decrypted_state_without_key_given() -> None: + create_fake_lean_cli_directory() + + project_path = Path.cwd() / "Python Project" + + encryption_file_path = project_path / "encryption_x.txt" + encryption_file_path.write_text("KtSwJtq5a4uuQmxbPqcCP3d8yMRz5TZxDBAKy7kGwPcvcvsNBdCprGYwSBN8ntJa5JNNYHTB2GrBpAbkA38kCdnceegffZH7") + + key_hash = get_project_key_hash(encryption_file_path) + keys_api_data = {'keys': [{'name': 'test', 'hash': key_hash}]} + + api_client = mock.Mock() + api_client.encryption_keys.list = mock.MagicMock(return_value=keys_api_data) + cloud_project = create_api_project(1, "Python Project") + api_client.projects.create = mock.MagicMock(return_value=cloud_project) + fake_cloud_files = [QCFullFile(name="file.py", content="testing", modified=datetime.now(), isLibrary=False)] + api_client.files.get_all = mock.MagicMock(return_value=fake_cloud_files) + # update the project with encrypted files so that it's in encrypted state + expected_encrypted_files = _get_expected_encrypted_files_content() + source_files = container.project_manager.get_source_files(project_path) + for source_file in source_files: + source_file.write_text(expected_encrypted_files[source_file.name]) + + 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)) + + result = CliRunner().invoke(lean, ["cloud", "push", "--project", project_path]) + + assert result.exit_code == 0 + source_files = container.project_manager.get_source_files(project_path) + expected_encrypted_files = _get_expected_encrypted_files_content() + for file in source_files: + assert expected_encrypted_files[file.name].strip() == file.read_text().strip() + expected_arguments = { + "name": "Python Project", + "description": "", + "files":[{'name': name, 'content': content} for name, content in expected_encrypted_files.items()], + "libraries": [], + "encryption_key": key_hash + } + api_client.projects.update.assert_called_once_with(1, **expected_arguments) + +def test_cloud_push_decrypted_when_local_files_in_decrypted_state_and_cloud_project_in_encrypted_state_without_key_given() -> None: + create_fake_lean_cli_directory() + + project_path = Path.cwd() / "Python Project" + + encryption_file_path = project_path / "encryption_x.txt" + encryption_file_path.write_text("KtSwJtq5a4uuQmxbPqcCP3d8yMRz5TZxDBAKy7kGwPcvcvsNBdCprGYwSBN8ntJa5JNNYHTB2GrBpAbkA38kCdnceegffZH7") + key_hash = get_project_key_hash(encryption_file_path) + + api_client = mock.Mock() + cloud_project = create_api_project(1, "Python Project", encrypted=True, encryptionKey={"name":"test", "id": key_hash}) + api_client.projects.create = mock.MagicMock(return_value=cloud_project) + fake_cloud_files = [QCFullFile(name="file.py", content="testing", modified=datetime.now(), isLibrary=False)] + api_client.files.get_all = mock.MagicMock(return_value=fake_cloud_files) + + init_container(api_client_to_use=api_client) + + result = CliRunner().invoke(lean, ["cloud", "push", "--project", project_path]) + + assert result.exit_code == 0 + expected_arguments = { + "name": "Python Project", + "description": "", + "files": [{'name': 'main.py', 'content': '# region imports\nfrom AlgorithmImports import *\n# endregion\n\nclass PythonProject(QCAlgorithm):\n\n def Initialize(self):\n # Locally Lean installs free sample data, to download more data please visit https://www.quantconnect.com/docs/v2/lean-cli/datasets/downloading-data\n self.SetStartDate(2013, 10, 7) # Set Start Date\n self.SetEndDate(2013, 10, 11) # Set End Date\n self.SetCash(100000) # Set Strategy Cash\n self.AddEquity("SPY", Resolution.Minute)\n\n def OnData(self, data: Slice):\n """OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.\n Arguments:\n data: Slice object keyed by symbol containing the stock data\n """\n if not self.Portfolio.Invested:\n self.SetHoldings("SPY", 1)\n self.Debug("Purchased Stock")\n'}, {'name': 'research.ipynb', 'content': '{\n "cells": [\n {\n "cell_type": "markdown",\n "metadata": {},\n "source": [\n "![QuantConnect Logo](https://cdn.quantconnect.com/web/i/icon.png)\\n",\n "
"\n ]\n },\n {\n "cell_type": "code",\n "execution_count": null,\n "metadata": {},\n "outputs": [],\n "source": [\n "# QuantBook Analysis Tool \\n",\n "# For more information see [https://www.quantconnect.com/docs/v2/our-platform/research/getting-started]\\n",\n "qb = QuantBook()\\n",\n "spy = qb.AddEquity(\\"SPY\\")\\n",\n "# Locally Lean installs free sample data, to download more data please visit https://www.quantconnect.com/docs/v2/lean-cli/datasets/downloading-data \\n",\n "qb.SetStartDate(2013, 10, 11)\\n",\n "history = qb.History(qb.Securities.Keys, 360, Resolution.Daily)\\n",\n "\\n",\n "# Indicator Analysis\\n",\n "bbdf = qb.Indicator(BollingerBands(30, 2), spy.Symbol, 360, Resolution.Daily)\\n",\n "bbdf.drop(\'standarddeviation\', axis=1).plot()"\n ]\n }\n ],\n "metadata": {\n "kernelspec": {\n "display_name": "Python 3",\n "language": "python",\n "name": "python3"\n }\n },\n "nbformat": 4,\n "nbformat_minor": 2\n}\n'}], + "libraries": [], + "encryption_key": '' + } + api_client.projects.update.assert_called_once_with(1, **expected_arguments) + +def test_cloud_push_aborts_when_local_files_in_encrypted_state_with_key_x_and_cloud_project_in_encrypted_state_with_key_y() -> None: + create_fake_lean_cli_directory() + + project_path = Path.cwd() / "Python Project" + + encryption_file_path_x = project_path / "encryption_x.txt" + encryption_file_path_x.write_text("KtSwJtq5a4uuQmxbPqcCP3d8yMRz5TZxDBAKy7kGwPcvcvsNBdCprGYwSBN8ntJa5JNNYHTB2GrBpAbkA38kCdnceegffZH7") + + encryption_file_path_y = project_path / "encryption_y.txt" + encryption_file_path_y.write_text("Jtq5a4uuQmxbPqcCP3d8yMRz5TZxDBAKy7kGwPcvcvsNBdCprGYwSBN8ntJa5JNNYHTB2GrBpAbkA38kCdnceegffZH7") + + # Keys API Data + key_hash_y = get_project_key_hash(encryption_file_path_y) + + api_client = mock.Mock() + cloud_project = create_api_project(1, "Python Project", encrypted=True, encryptionKey={"name":"test", "id": key_hash_y}) + api_client.projects.create = mock.MagicMock(return_value=cloud_project, encrypted=True, encryptionKey={"name":"test", "id": key_hash_y}) + fake_cloud_files = [QCFullFile(name="file.py", content="testing", modified=datetime.now(), isLibrary=False)] + 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_x)) + + result = CliRunner().invoke(lean, ["cloud", "push", "--project", project_path]) assert result.exit_code == 0 + api_client.projects.update.assert_not_called() + + +def _get_expected_encrypted_files_content() -> dict: + return { + "main.py": + """UpMdqgoXS1tgqGgy6nKkmlxrOV7ikoc5oJAmS+pcMcmD0qsfJq5GE/yvdg9mucXrhfgLjD7of3YalHYC +mJcVeO1VSlHCA3oNq5kS82YV4Rt0KL0IApPXlV7yAvJW/SbqcziwPXb/OSmHroQfaD9IJWJdNXrcjfs1 +CAXq8mV0BX0OipDTnQL2EfHCDNn3PL7lJPg0lYOKjSuMf6QgvFWElT6Kw0UDiLtKlU6jJwsKxeJDDRAw +KkMJVrrIsxk5g92ERCquoTmYHVm1lf0xdsM+vTqSRxwucEfaoi7DgA9SkJKgaVoDgx90ScqssPzdlsgJ +zIUdPbHO1GN1D5ZfHXU+7Sww4pHDh6QTyLYnuOHnPQ8IuXYDWQPdw9QaHrbeuWqqwyFl859Tra6RZFlD +136K2CZ6PcClbGeeGO2NmGPRFuwHH5B5aSfUs8Y3OMdYnZylXtuWAMPU6xEZIWfGHIOZlQZ7s06rvfYT +CfMHdkFCRi9QnI1vRbqO4yhBDB/BbAG9W7wt/FX1f+1hHk3xpaP+MLufoN48xXSxudIV975Jxdoic8YK +P/6UZ5FWDVTZDRzJ7e9hiS5LOr410G8PqDOr0JuNap3t2FvGLuuR1CFfXsHm2QcsaJgiODzAa/xSW2vq +ffQPkZxz3bhGpuJXoWe18B92m0/Scih5J224QMXoZJYs6UYEl6eQX67mlGSSntotykByWB7upnPFDPSU +l9if+gA+GrW3hyVhDSY/5W9mUA6mw7lelE0b+0cAfuEczDXInQeMbQjAtEOJufIQl7NSZVzc4BWOFoXS +VK3OnsQnbg2f9WN7mQ4UE2xVoVAOudDPSXh21HjhUa9DbqrFy4FQN66G8LgbF6Sw58d+HfCIPOVoNl1f +mMJBPythF0UjBYwle+2EpDC/9WdsVrm02OCEehbs4gfG0vOTtmwMY0Y4Urpxs01ATOfoZEg5gRPMFRD7 +zFhH4MGKrPdBA3NXnvRozJsdHRBEVPhmuhtideHMKs68o8BwDLJ8oY6dyyxpjlFhGlHELGDTo7ljCRtR +6d5M+DpJn9vEZ9xWFe60vkgIq3dpn3SwPnrrsd2Ee5TwagDF2Bj8EhMVhshlxjrNR9NNguYgZVFUAJEM +InB4pibT1Q0FLnwFOp84t8offZ5w+hmGm7Jggs99qmx2Tmu4u4067A==""" +, +"research.ipynb": + """NIiAgzU8gzaJ3YEIWysBh0e0xxm9rpAWDE4Pir/wzKtla/wcbs98GU5cdOxgd7Gjlcu0zNbFzE5x8Eyx +qSuh3cQU17xQSisPxvjfDD/h2z9AnFT1jD1Vhc+Nn5ngwpgCA6P5fHT4VhPgaKDp7r9zc8pAURcSd04M +2dmHtGs89it3QFrYNgYAvN4kIYVZuROhnUSiYN+y7kzWzLKLIK3a+y6R4ibr9ju5S0DCtIS87MwrHFbi +NmS0mqnlYmBqLPVijrgAJYu4YmqrOb+LxwQSUXL80UsgddUFtgKWDKGyTVFs0o92/x9OT2XTCcMOhhQg +X/6h5c99rP5W3GqkcDaaiueShSD43u3LXnijK6yugUqALcMIBLxF4Iczq3xLov9MfrdgDUaxFEOu0dDK +JpecJeSx9TxVUxzcWvn3alUycScBRV1w8VH0Wa4Lf9vc1YIOe9ITy/hLW2+QEBAHxrhYd7vq1MVVDdRI +lbXtLVMOVI9YYBNEunf3drQm6fCnKv+arnt/rYSIQDvhrlqem6qKS3JE4V8gVzXQDefTewx3/hppDOxL ++ccikaUcVkIAULF1Jwp9qxCpCsiz5vLXynD47pf3mhS2FC6Dd1g30xXUVNeTmpRE7TfS/gHzkyDrtTvC +LBvBnkImsJQjDCJl/NzLlFRh4wiY4SL/bTdxL2YZJebJ1zdw3PXQhvWvgztG0wmIRE/U+1xv95gc5w5x +7N2WT7F9KVHI6NrtrcU6c5hWo4q/QEBxb5ETSyNpgsVHvKxKblPGslki/aH6WOexu7tSzbuA9rSwVgiS +NWh9y3qZ6aMOMq3dDAr4wGBkkGQsirasCEt3YEa9rG5CDH8JYApCeZk39zt+oqjQ4DcYprgW/q87fFRB +thxY4cM3BGG/IIpX9oAxvFNe7PSXIx3BBkyhqiJqJRhS3YfvasBWFnJd11h9JXVsKJudVBoJvwrzxldU +Td2mSr1d//xkIe6TeqPK3ivofMoCNoDYyO57juki8ban9Z8q0lpgShWBM8zUPNnezctwJZIp6aVaVdfC +nZtG49f06w1IRC0dOj6yM+adiEPx/vKxrCn1sqd0eoAeKIs5iEm398YJlt9G0eubSSpoDC6iH1GjwH8A +NFNRoKncB4c3ocmajgNda4o32UfpwvNr+fKLmpP+57dulpA09i/FqwF2L8dDNm4JSTjfNdeooM9TIDsG +b3CIItqmNnve8fryEykjE8WgJaP70rnsXRY6WHLivy7OTNMKG8Ij9MOyx44Wq3TWY/omRDjYgfCO4QRZ +vTtlnIMp8VtLlBgxRc6TBPZutYZXDJyY4N4n2Jpspl3KUbe3iHNm6nwl0SCwrIl3aYbCs55BP/LOGXLx +dNIZzkxJhnfLtomrxO4QgcUnjNEWYJ1GvRAplTa99RFFYXOuMQID4+UH+6RM2e6cqBafv+2+lpetfbBM +gMOA+qNdar9Y+b+vLuSz2ak0MCn3h8cAAlXgKNpo4shbgn+R4TViLGfhUZbcSG9VZ0UEBYyg5nxXggqh +fPkc80yGC7lwo0XmrQrp62IaQy3iB2MYWB29KCX9MdghcobfhmW3BEtYqDRLE61/MWYi9scowMkf0TxT +DyMRNS9qVwJ+2gS44xD+lQmQuE1aM+rBmRkVKWA+fB5HzZe6yyN8GkM4D23hcw+DBs9opXvq6QyGd8Le +KqLYkbp7Be44+d6mie1aGp0hRUjF9LzNapPdnYusx4QwoIlt0wYTr1waqAYMG73PhDI36B6oa4tOwuSj +86Wje1jEfUK01LVuVYloclnG3fTQxQlUcrRyQNCbrg0M5elYWbjYybx51/gb//9McGYXW0SWltNlvLvL +SUroK4+Yhqfk5MVeascvM94fFBfieT3erU8D+iCiaVvZGcuE5mE0bLBquXpMuO4fQ06XVI/NzlhVILrV +yXa/cl9WaYDFPqM+eFVJfg==""" +} - project_config.set.assert_called_with("organization-id", "123") diff --git a/tests/commands/test_decrypt.py b/tests/commands/test_decrypt.py new file mode 100644 index 00000000..a101db4b --- /dev/null +++ b/tests/commands/test_decrypt.py @@ -0,0 +1,207 @@ +# QUANTCONNECT.COM - Democratizing Finance, Empowering Individuals. +# Lean CLI v1.0. Copyright 2021 QuantConnect Corporation. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from pathlib import Path +from click.testing import CliRunner + +from lean.commands import lean +from tests.test_helpers import create_fake_lean_cli_directory +from lean.container import container + + +def test_decrypt_decrypts_file_in_case_project_not_in_decrypt_state() -> None: + create_fake_lean_cli_directory() + + project_path = Path.cwd() / "Python Project" + + source_files = container.project_manager.get_source_files(project_path) + file_contents_map = {file.name: file.read_text() for file in source_files} + + encryption_file_path = project_path / "encryption_x.txt" + encryption_file_path.write_text("KtSwJtq5a4uuQmxbPqcCP3d8yMRz5TZxDBAKy7kGwPcvcvsNBdCprGYwSBN8ntJa5JNNYHTB2GrBpAbkA38kCdnceegffZH7") + + result = CliRunner().invoke(lean, ["encrypt", "Python Project", "--key", encryption_file_path]) + + 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) != None + for file in source_files: + assert file_contents_map[file.name] != file.read_text() + + result = CliRunner().invoke(lean, ["decrypt", "Python Project", "--key", encryption_file_path]) + + assert result.exit_code == 0 + + source_files = container.project_manager.get_source_files(project_path) + for file in source_files: + assert file_contents_map[file.name] == file.read_text() + 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) == None + + +def test_decrypt_does_not_change_file_in_case_project_already_in_decrypt_state() -> None: + create_fake_lean_cli_directory() + + project_path = Path.cwd() / "Python Project" + + 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) == None + + source_files = container.project_manager.get_source_files(project_path) + file_contents_map = {file.name: file.read_text() for file in source_files} + encryption_file_path = project_path / "encryption_x.txt" + encryption_file_path.write_text("KtSwJtq5a4uuQmxbPqcCP3d8yMRz5TZxDBAKy7kGwPcvcvsNBdCprGYwSBN8ntJa5JNNYHTB2GrBpAbkA38kCdnceegffZH7") + + result = CliRunner().invoke(lean, ["decrypt", "Python Project", "--key", encryption_file_path]) + + assert result.exit_code == 0 + + source_files = container.project_manager.get_source_files(project_path) + for file in source_files: + assert file_contents_map[file.name] == file.read_text() + 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) == None + + +def test_decrypt_uses_key_from_config_file_when_not_provided() -> None: + create_fake_lean_cli_directory() + + project_path = Path.cwd() / "Python Project" + + source_files = container.project_manager.get_source_files(project_path) + file_contents_map = {file.name: file.read_text() for file in source_files} + + encryption_file_path = project_path / "encryption_x.txt" + encryption_file_path.write_text("KtSwJtq5a4uuQmxbPqcCP3d8yMRz5TZxDBAKy7kGwPcvcvsNBdCprGYwSBN8ntJa5JNNYHTB2GrBpAbkA38kCdnceegffZH7") + + CliRunner().invoke(lean, ["encrypt", "Python Project", "--key", encryption_file_path]) + + 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) != None + for file in source_files: + assert file_contents_map[file.name] != file.read_text() + + result = CliRunner().invoke(lean, ["decrypt", "Python Project"]) + + assert result.exit_code == 0 + + source_files = container.project_manager.get_source_files(project_path) + for file in source_files: + assert file_contents_map[file.name] == file.read_text() + 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) == None + +def test_decrypt_updates_project_config_file() -> None: + create_fake_lean_cli_directory() + + project_path = Path.cwd() / "Python Project" + + source_files = container.project_manager.get_source_files(project_path) + file_contents_map = {file.name: file.read_text() for file in source_files} + + encryption_file_path = project_path / "encryption_x.txt" + encryption_file_path.write_text("KtSwJtq5a4uuQmxbPqcCP3d8yMRz5TZxDBAKy7kGwPcvcvsNBdCprGYwSBN8ntJa5JNNYHTB2GrBpAbkA38kCdnceegffZH7") + + CliRunner().invoke(lean, ["encrypt", "Python Project", "--key", encryption_file_path]) + + 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) != None + for file in source_files: + assert file_contents_map[file.name] != file.read_text() + + result = CliRunner().invoke(lean, ["decrypt", "Python Project", "--key", encryption_file_path]) + + assert result.exit_code == 0 + + source_files = container.project_manager.get_source_files(project_path) + for file in source_files: + assert file_contents_map[file.name] == file.read_text() + 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) == None + +def test_decrypt_does_not_update_project_config_file_if_not_all_files_successful() -> None: + create_fake_lean_cli_directory() + + project_path = Path.cwd() / "Python Project" + + encryption_file_path = project_path / "encryption_x.txt" + encryption_file_path.write_text("KtSwJtq5a4uuQmxbPqcCP3d8yMRz5TZxDBAKy7kGwPcvcvsNBdCprGYwSBN8ntJa5JNNYHTB2GrBpAbkA38kCdnceegffZH7") + + CliRunner().invoke(lean, ["encrypt", "Python Project", "--key", encryption_file_path]) + + 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) != None + + # let's corrupt one file + source_files = container.project_manager.get_source_files(project_path) + source_files[0].write_text("corrupted") + file_contents_map = {file.name: file.read_text() for file in source_files} + + result = CliRunner().invoke(lean, ["decrypt", "Python Project", "--key", encryption_file_path]) + + assert result.exit_code != 0 + + source_files = container.project_manager.get_source_files(project_path) + for file in source_files: + assert file_contents_map[file.name] == file.read_text() + 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) != None + +def test_decrypt_aborts_when_key_is_not_provided_and_not_in_config_file() -> None: + create_fake_lean_cli_directory() + + project_path = Path.cwd() / "Python Project" + + project_config = container.project_config_manager.get_project_config(project_path) + source_files = container.project_manager.get_source_files(project_path) + file_contents_map = {file.name: file.read_text() for file in source_files} + assert project_config.get("encrypted", False) == False + project_config.set("encrypted", True) + + result = CliRunner().invoke(lean, ["decrypt", "Python Project"]) + + assert result.exit_code != 0 + + source_files = container.project_manager.get_source_files(project_path) + for file in source_files: + assert file_contents_map[file.name] == file.read_text() + +def test_decrypt_aborts_when_provided_key_different_from_key_in_config_file() -> None: + create_fake_lean_cli_directory() + + project_path = Path.cwd() / "Python Project" + + encryption_file_path_x = project_path / "encryption_x.txt" + encryption_file_path_x.write_text("KtSwJtq5a4uuQmxbPqcCP3d8yMRz5TZxDBAKy7kGwPcvcvsNBdCprGYwSBN8ntJa5JNNYHTB2GrBpAbkA38kCdnceegffZH7") + + encryption_file_path_y = project_path / "encryption_y.txt" + encryption_file_path_y.write_text("Jtq5a4uuQmxbPqcCP3d8yMRz5TZxDBAKy7kGwPcvcvsNBdCprGYwSBN8ntJa5JNNYHTB2GrBpAbkA38kCdnceegffZH7") + + CliRunner().invoke(lean, ["encrypt", "Python Project", "--key", encryption_file_path_x]) + + 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) != None + + result = CliRunner().invoke(lean, ["decrypt", "Python Project", "--key", encryption_file_path_y]) + + assert result.exit_code != 0 diff --git a/tests/commands/test_encrypt.py b/tests/commands/test_encrypt.py new file mode 100644 index 00000000..1b87b82e --- /dev/null +++ b/tests/commands/test_encrypt.py @@ -0,0 +1,196 @@ +# QUANTCONNECT.COM - Democratizing Finance, Empowering Individuals. +# Lean CLI v1.0. Copyright 2021 QuantConnect Corporation. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from pathlib import Path +from click.testing import CliRunner + +from lean.commands import lean +from tests.test_helpers import create_fake_lean_cli_directory +from lean.container import container + + +def test_encrypt_encrypts_file_in_case_project_not_in_encrypt_state() -> None: + create_fake_lean_cli_directory() + + project_path = Path.cwd() / "Python Project" + + encryption_file_path = project_path / "encryption_x.txt" + encryption_file_path.write_text("KtSwJtq5a4uuQmxbPqcCP3d8yMRz5TZxDBAKy7kGwPcvcvsNBdCprGYwSBN8ntJa5JNNYHTB2GrBpAbkA38kCdnceegffZH7") + + 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) == None + + result = CliRunner().invoke(lean, ["encrypt", "Python Project", "--key", encryption_file_path]) + + assert result.exit_code == 0 + + source_files = container.project_manager.get_source_files(project_path) + expected_encrypted_files = _get_expected_encrypted_files_content() + for file in source_files: + assert expected_encrypted_files[file.name].strip() == file.read_text() + 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) + + +def test_encrypt_does_not_change_file_in_case_project_already_in_encrypt_state() -> None: + create_fake_lean_cli_directory() + + project_path = Path.cwd() / "Python Project" + + encryption_file_path = project_path / "encryption_x.txt" + encryption_file_path.write_text("KtSwJtq5a4uuQmxbPqcCP3d8yMRz5TZxDBAKy7kGwPcvcvsNBdCprGYwSBN8ntJa5JNNYHTB2GrBpAbkA38kCdnceegffZH7") + + 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)) + + source_files = container.project_manager.get_source_files(project_path) + file_contents_map = {file.name: file.read_text() for file in source_files} + + result = CliRunner().invoke(lean, ["encrypt", "Python Project", "--key", encryption_file_path]) + + assert result.exit_code == 0 + + source_files = container.project_manager.get_source_files(project_path) + for file in source_files: + assert file_contents_map[file.name] == file.read_text() + 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) + + +def test_encrypt_uses_key_from_config_file_when_not_provided() -> None: + create_fake_lean_cli_directory() + + project_path = Path.cwd() / "Python Project" + + encryption_file_path = project_path / "encryption_x.txt" + encryption_file_path.write_text("KtSwJtq5a4uuQmxbPqcCP3d8yMRz5TZxDBAKy7kGwPcvcvsNBdCprGYwSBN8ntJa5JNNYHTB2GrBpAbkA38kCdnceegffZH7") + project_config = container.project_config_manager.get_project_config(project_path) + project_config.set("encryption-key-path", str(encryption_file_path)) + + result = CliRunner().invoke(lean, ["encrypt", "Python Project"]) + + assert result.exit_code == 0 + + source_files = container.project_manager.get_source_files(project_path) + expected_encrypted_files = _get_expected_encrypted_files_content() + for file in source_files: + assert expected_encrypted_files[file.name].strip() == file.read_text() + 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) + +def test_encrypt_updates_project_config_file() -> None: + create_fake_lean_cli_directory() + + project_path = Path.cwd() / "Python Project" + + encryption_file_path = project_path / "encryption_x.txt" + encryption_file_path.write_text("KtSwJtq5a4uuQmxbPqcCP3d8yMRz5TZxDBAKy7kGwPcvcvsNBdCprGYwSBN8ntJa5JNNYHTB2GrBpAbkA38kCdnceegffZH7") + + result = CliRunner().invoke(lean, ["encrypt", "Python Project", "--key", encryption_file_path]) + + assert result.exit_code == 0 + + source_files = container.project_manager.get_source_files(project_path) + expected_encrypted_files = _get_expected_encrypted_files_content() + for file in source_files: + assert expected_encrypted_files[file.name].strip() == file.read_text() + 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) + +def test_encrypt_aborts_when_key_is_not_provided_and_not_in_config_file() -> None: + create_fake_lean_cli_directory() + + project_path = Path.cwd() / "Python Project" + + source_files = container.project_manager.get_source_files(project_path) + file_contents_map = {file.name: file.read_text() for file in source_files} + + result = CliRunner().invoke(lean, ["encrypt", "Python Project"]) + + assert result.exit_code != 0 + + source_files = container.project_manager.get_source_files(project_path) + for file in source_files: + assert file_contents_map[file.name] == file.read_text() + +def test_encrypt_aborts_when_provided_key_different_from_key_in_config_file() -> None: + create_fake_lean_cli_directory() + + project_path = Path.cwd() / "Python Project" + + encryption_file_path_x = project_path / "encryption_x.txt" + encryption_file_path_x.write_text("KtSwJtq5a4uuQmxbPqcCP3d8yMRz5TZxDBAKy7kGwPcvcvsNBdCprGYwSBN8ntJa5JNNYHTB2GrBpAbkA38kCdnceegffZH7") + + encryption_file_path_y = project_path / "encryption_y.txt" + encryption_file_path_y.write_text("Jtq5a4uuQmxbPqcCP3d8yMRz5TZxDBAKy7kGwPcvcvsNBdCprGYwSBN8ntJa5JNNYHTB2GrBpAbkA38kCdnceegffZH7") + + project_config = container.project_config_manager.get_project_config(project_path) + project_config.set("encryption-key-path", str(encryption_file_path_x)) + + result = CliRunner().invoke(lean, ["encrypt", "Python Project", "--key", encryption_file_path_y]) + + assert result.exit_code != 0 + + +def _get_expected_encrypted_files_content() -> dict: + return { + "main.py": + """UpMdqgoXS1tgqGgy6nKkmlxrOV7ikoc5oJAmS+pcMcmD0qsfJq5GE/yvdg9mucXrhfgLjD7of3YalHYC +mJcVeO1VSlHCA3oNq5kS82YV4Rt0KL0IApPXlV7yAvJW/SbqcziwPXb/OSmHroQfaD9IJWJdNXrcjfs1 +CAXq8mV0BX0OipDTnQL2EfHCDNn3PL7lJPg0lYOKjSuMf6QgvFWElT6Kw0UDiLtKlU6jJwsKxeJDDRAw +KkMJVrrIsxk5g92ERCquoTmYHVm1lf0xdsM+vTqSRxwucEfaoi7DgA9SkJKgaVoDgx90ScqssPzdlsgJ +zIUdPbHO1GN1D5ZfHXU+7Sww4pHDh6QTyLYnuOHnPQ8IuXYDWQPdw9QaHrbeuWqqwyFl859Tra6RZFlD +136K2CZ6PcClbGeeGO2NmGPRFuwHH5B5aSfUs8Y3OMdYnZylXtuWAMPU6xEZIWfGHIOZlQZ7s06rvfYT +CfMHdkFCRi9QnI1vRbqO4yhBDB/BbAG9W7wt/FX1f+1hHk3xpaP+MLufoN48xXSxudIV975Jxdoic8YK +P/6UZ5FWDVTZDRzJ7e9hiS5LOr410G8PqDOr0JuNap3t2FvGLuuR1CFfXsHm2QcsaJgiODzAa/xSW2vq +ffQPkZxz3bhGpuJXoWe18B92m0/Scih5J224QMXoZJYs6UYEl6eQX67mlGSSntotykByWB7upnPFDPSU +l9if+gA+GrW3hyVhDSY/5W9mUA6mw7lelE0b+0cAfuEczDXInQeMbQjAtEOJufIQl7NSZVzc4BWOFoXS +VK3OnsQnbg2f9WN7mQ4UE2xVoVAOudDPSXh21HjhUa9DbqrFy4FQN66G8LgbF6Sw58d+HfCIPOVoNl1f +mMJBPythF0UjBYwle+2EpDC/9WdsVrm02OCEehbs4gfG0vOTtmwMY0Y4Urpxs01ATOfoZEg5gRPMFRD7 +zFhH4MGKrPdBA3NXnvRozJsdHRBEVPhmuhtideHMKs68o8BwDLJ8oY6dyyxpjlFhGlHELGDTo7ljCRtR +6d5M+DpJn9vEZ9xWFe60vkgIq3dpn3SwPnrrsd2Ee5TwagDF2Bj8EhMVhshlxjrNR9NNguYgZVFUAJEM +InB4pibT1Q0FLnwFOp84t8offZ5w+hmGm7Jggs99qmx2Tmu4u4067A==""" +, +"research.ipynb": + """NIiAgzU8gzaJ3YEIWysBh0e0xxm9rpAWDE4Pir/wzKtla/wcbs98GU5cdOxgd7Gjlcu0zNbFzE5x8Eyx +qSuh3cQU17xQSisPxvjfDD/h2z9AnFT1jD1Vhc+Nn5ngwpgCA6P5fHT4VhPgaKDp7r9zc8pAURcSd04M +2dmHtGs89it3QFrYNgYAvN4kIYVZuROhnUSiYN+y7kzWzLKLIK3a+y6R4ibr9ju5S0DCtIS87MwrHFbi +NmS0mqnlYmBqLPVijrgAJYu4YmqrOb+LxwQSUXL80UsgddUFtgKWDKGyTVFs0o92/x9OT2XTCcMOhhQg +X/6h5c99rP5W3GqkcDaaiueShSD43u3LXnijK6yugUqALcMIBLxF4Iczq3xLov9MfrdgDUaxFEOu0dDK +JpecJeSx9TxVUxzcWvn3alUycScBRV1w8VH0Wa4Lf9vc1YIOe9ITy/hLW2+QEBAHxrhYd7vq1MVVDdRI +lbXtLVMOVI9YYBNEunf3drQm6fCnKv+arnt/rYSIQDvhrlqem6qKS3JE4V8gVzXQDefTewx3/hppDOxL ++ccikaUcVkIAULF1Jwp9qxCpCsiz5vLXynD47pf3mhS2FC6Dd1g30xXUVNeTmpRE7TfS/gHzkyDrtTvC +LBvBnkImsJQjDCJl/NzLlFRh4wiY4SL/bTdxL2YZJebJ1zdw3PXQhvWvgztG0wmIRE/U+1xv95gc5w5x +7N2WT7F9KVHI6NrtrcU6c5hWo4q/QEBxb5ETSyNpgsVHvKxKblPGslki/aH6WOexu7tSzbuA9rSwVgiS +NWh9y3qZ6aMOMq3dDAr4wGBkkGQsirasCEt3YEa9rG5CDH8JYApCeZk39zt+oqjQ4DcYprgW/q87fFRB +thxY4cM3BGG/IIpX9oAxvFNe7PSXIx3BBkyhqiJqJRhS3YfvasBWFnJd11h9JXVsKJudVBoJvwrzxldU +Td2mSr1d//xkIe6TeqPK3ivofMoCNoDYyO57juki8ban9Z8q0lpgShWBM8zUPNnezctwJZIp6aVaVdfC +nZtG49f06w1IRC0dOj6yM+adiEPx/vKxrCn1sqd0eoAeKIs5iEm398YJlt9G0eubSSpoDC6iH1GjwH8A +NFNRoKncB4c3ocmajgNda4o32UfpwvNr+fKLmpP+57dulpA09i/FqwF2L8dDNm4JSTjfNdeooM9TIDsG +b3CIItqmNnve8fryEykjE8WgJaP70rnsXRY6WHLivy7OTNMKG8Ij9MOyx44Wq3TWY/omRDjYgfCO4QRZ +vTtlnIMp8VtLlBgxRc6TBPZutYZXDJyY4N4n2Jpspl3KUbe3iHNm6nwl0SCwrIl3aYbCs55BP/LOGXLx +dNIZzkxJhnfLtomrxO4QgcUnjNEWYJ1GvRAplTa99RFFYXOuMQID4+UH+6RM2e6cqBafv+2+lpetfbBM +gMOA+qNdar9Y+b+vLuSz2ak0MCn3h8cAAlXgKNpo4shbgn+R4TViLGfhUZbcSG9VZ0UEBYyg5nxXggqh +fPkc80yGC7lwo0XmrQrp62IaQy3iB2MYWB29KCX9MdghcobfhmW3BEtYqDRLE61/MWYi9scowMkf0TxT +DyMRNS9qVwJ+2gS44xD+lQmQuE1aM+rBmRkVKWA+fB5HzZe6yyN8GkM4D23hcw+DBs9opXvq6QyGd8Le +KqLYkbp7Be44+d6mie1aGp0hRUjF9LzNapPdnYusx4QwoIlt0wYTr1waqAYMG73PhDI36B6oa4tOwuSj +86Wje1jEfUK01LVuVYloclnG3fTQxQlUcrRyQNCbrg0M5elYWbjYybx51/gb//9McGYXW0SWltNlvLvL +SUroK4+Yhqfk5MVeascvM94fFBfieT3erU8D+iCiaVvZGcuE5mE0bLBquXpMuO4fQ06XVI/NzlhVILrV +yXa/cl9WaYDFPqM+eFVJfg==""" +} diff --git a/tests/components/util/test_pull_manager.py b/tests/components/util/test_pull_manager.py index 42e78a57..82ed3d92 100644 --- a/tests/components/util/test_pull_manager.py +++ b/tests/components/util/test_pull_manager.py @@ -32,7 +32,9 @@ def _create_pull_manager(api_client: mock.Mock, logger = mock.Mock() platform_manager = mock.Mock() project_manager = container.project_manager - return PullManager(logger, api_client, project_manager, project_config_manager, library_manager, platform_manager) + organization_manager = container.organization_manager + return PullManager(logger, api_client, project_manager, + project_config_manager, library_manager, platform_manager, organization_manager) def _make_cloud_projects_and_libraries(project_count: int, @@ -65,8 +67,13 @@ def _assert_pull_manager_adds_property_to_project_config(prop: str, api_client.lean.environments = mock.MagicMock(return_value=create_lean_environments()) api_client.files.get_all = mock.MagicMock(return_value=[]) + def config_get_side_effect(key: str, default: Any = None): + if key == "encryption-key-path": + return None + return [] + project_config = mock.Mock() - project_config.get = mock.MagicMock(return_value=[]) + project_config.get = mock.MagicMock(side_effect=config_get_side_effect) project_config.set = mock.Mock() project_config_manager = mock.Mock() @@ -75,7 +82,7 @@ def _assert_pull_manager_adds_property_to_project_config(prop: str, pull_manager = _create_pull_manager(api_client, project_config_manager) pull_manager.pull_projects(cloud_projects, cloud_projects) - project_config.set.assert_called_with(prop, expected_value) + project_config.set.assert_any_call(prop, expected_value) def test_pull_manager_adds_lean_engine_version_to_config() -> None: @@ -107,8 +114,13 @@ def _assert_pull_manager_removes_property_from_project_config(prop: str, cloud_p api_client.lean.environments = mock.MagicMock(return_value=create_lean_environments()) api_client.files.get_all = mock.MagicMock(return_value=[]) + def config_get_side_effect(key: str, default: Any = None): + if key == "encryption-key-path": + return None + return [] + project_config = mock.Mock() - project_config.get = mock.MagicMock(return_value=[]) + project_config.get = mock.MagicMock(side_effect=config_get_side_effect) project_config.set = mock.Mock() project_config.delete = mock.Mock() @@ -316,8 +328,13 @@ def test_pull_projects_updates_lean_config() -> None: api_client.projects.get_all.return_value = cloud_projects api_client.files.get_all = mock.MagicMock(return_value=[]) + def config_get_side_effect(key: str, default: Any = None): + if key == "encryption-key-path": + return None + return [] + project_config = mock.Mock() - project_config.get = mock.MagicMock(return_value=[]) # get("libraries") + project_config.get = mock.MagicMock(side_effect=config_get_side_effect) project_config_manager = mock.Mock() project_config_manager.get_project_config = mock.MagicMock(return_value=project_config) diff --git a/tests/test_helpers.py b/tests/test_helpers.py index 3ce821c5..7c979d98 100644 --- a/tests/test_helpers.py +++ b/tests/test_helpers.py @@ -19,7 +19,7 @@ from lean.commands.create_project import (DEFAULT_CSHARP_MAIN, DEFAULT_CSHARP_NOTEBOOK, DEFAULT_PYTHON_MAIN, DEFAULT_PYTHON_NOTEBOOK, LIBRARY_PYTHON_MAIN, LIBRARY_CSHARP_MAIN) from lean.components.util.project_manager import ProjectManager -from lean.models.api import QCLanguage, QCProject, QCFullOrganization, \ +from lean.models.api import QCLanguage, QCProject, QCFullOrganization, ProjectEncryptionKey, \ QCOrganizationData, QCOrganizationCredit, QCNode, QCNodeList, QCNodePrice, QCLeanEnvironment @@ -136,7 +136,7 @@ def create_fake_lean_cli_directory_with_subdirectories(depth: int) -> None: _write_fake_directory(files) -def create_api_project(id: int, name: str) -> QCProject: +def create_api_project(id: int, name: str, encrypted: bool=False, encryptionKey: ProjectEncryptionKey=None) -> QCProject: """Creates a fake API project response.""" return QCProject( projectId=id, @@ -151,7 +151,9 @@ def create_api_project(id: int, name: str) -> QCProject: leanPinnedToMaster=True, leanEnvironment=1, parameters=[], - libraries=[] + libraries=[], + encrypted= encrypted, + encryptionKey=encryptionKey )