From ba4b2c020a890015aca348880c5cbe46984ebe2c Mon Sep 17 00:00:00 2001 From: Arnav Sacheti Date: Tue, 26 Jul 2022 13:59:08 -0700 Subject: [PATCH 1/4] Creating KMS gateway and service (#366) Summary: Pull Request resolved: https://github.com/facebookresearch/fbpcp/pull/366 # What are Grant Tokens https://docs.aws.amazon.com/kms/latest/developerguide/grants.html Reviewed By: ziqih Differential Revision: D37400919 fbshipit-source-id: 28d9a3ed111b5a1a46ff59fbc575fb9b5e55313c --- fbpcp/gateway/kms.py | 27 +++++++++++++++++++++++ fbpcp/service/key_management_aws.py | 33 +++++++++++++++++++++++++++++ 2 files changed, 60 insertions(+) create mode 100644 fbpcp/gateway/kms.py create mode 100644 fbpcp/service/key_management_aws.py diff --git a/fbpcp/gateway/kms.py b/fbpcp/gateway/kms.py new file mode 100644 index 00000000..a78dc19d --- /dev/null +++ b/fbpcp/gateway/kms.py @@ -0,0 +1,27 @@ +#!/usr/bin/env python3 +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +from typing import Any, Dict, Optional + +import boto3 +from botocore.client import BaseClient +from fbpcp.gateway.aws import AWSGateway + + +class KMSGateway(AWSGateway): + def __init__( + self, + region: str, + access_key_id: Optional[str] = None, + access_key_data: Optional[str] = None, + config: Optional[Dict[str, Any]] = None, + ) -> None: + super().__init__(region, access_key_id, access_key_data, config) + self.client: BaseClient = boto3.client( + "kms", region_name=self.region, **self.config + ) diff --git a/fbpcp/service/key_management_aws.py b/fbpcp/service/key_management_aws.py new file mode 100644 index 00000000..dd2f8af6 --- /dev/null +++ b/fbpcp/service/key_management_aws.py @@ -0,0 +1,33 @@ +#!/usr/bin/env python3 +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +from typing import Any, Dict, List, Optional + +from fbpcp.gateway.kms import KMSGateway + + +class AWSKeyManagementService: + key_id: str + grant_tokens: List[str] + + def __init__( + self, + region: str, + key_id: str, + grant_tokens: Optional[List[str]] = None, + access_key_id: Optional[str] = None, + access_key_data: Optional[str] = None, + config: Optional[Dict[str, Any]] = None, + ) -> None: + """ + Args: + grant_tokens: Advertiser side specific, allows anyone with a grant token to have permisions to certain functions on KMS (Admin Controlled) + """ + self.kms_gateway = KMSGateway(region, access_key_id, access_key_data, config) + self.key_id = key_id + self.grant_tokens = grant_tokens if grant_tokens else [] From 7a767611c7eb15c75a0ece917f8955b0456b7ff6 Mon Sep 17 00:00:00 2001 From: Arnav Sacheti Date: Tue, 26 Jul 2022 13:59:08 -0700 Subject: [PATCH 2/4] Adding support for message signature (#377) Summary: Pull Request resolved: https://github.com/facebookresearch/fbpcp/pull/377 Reviewed By: ziqih Differential Revision: D37974927 fbshipit-source-id: 410aa1b31bfddcd9db35dbede0b7e75180345f79 --- fbpcp/gateway/kms.py | 23 +++++++++++- fbpcp/service/key_management.py | 15 ++++++++ fbpcp/service/key_management_aws.py | 18 +++++++++- tests/gateway/test_kms.py | 47 +++++++++++++++++++++++++ tests/service/test_key_managment_aws.py | 47 +++++++++++++++++++++++++ 5 files changed, 148 insertions(+), 2 deletions(-) create mode 100644 fbpcp/service/key_management.py create mode 100644 tests/gateway/test_kms.py create mode 100644 tests/service/test_key_managment_aws.py diff --git a/fbpcp/gateway/kms.py b/fbpcp/gateway/kms.py index a78dc19d..40aefcb3 100644 --- a/fbpcp/gateway/kms.py +++ b/fbpcp/gateway/kms.py @@ -6,10 +6,12 @@ # pyre-strict -from typing import Any, Dict, Optional +from base64 import b64encode +from typing import Any, Dict, List, Optional import boto3 from botocore.client import BaseClient +from fbpcp.decorator.error_handler import error_handler from fbpcp.gateway.aws import AWSGateway @@ -25,3 +27,22 @@ def __init__( self.client: BaseClient = boto3.client( "kms", region_name=self.region, **self.config ) + + @error_handler + def sign( + self, + key_id: str, + message: bytes, + message_type: str, + grant_tokens: List[str], + signing_algorithm: str, + ) -> str: + response = self.client.sign( + KeyId=key_id, + Message=message, + MessageType=message_type, + GrantTokens=grant_tokens, + SigningAlgorithm=signing_algorithm, + ) + signature = b64encode(response["Signature"]).decode() + return signature diff --git a/fbpcp/service/key_management.py b/fbpcp/service/key_management.py new file mode 100644 index 00000000..e9a83235 --- /dev/null +++ b/fbpcp/service/key_management.py @@ -0,0 +1,15 @@ +#!/usr/bin/env python3 +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +import abc + + +class KeyManagementService(abc.ABC): + @abc.abstractmethod + def sign(self, message: str) -> str: + pass diff --git a/fbpcp/service/key_management_aws.py b/fbpcp/service/key_management_aws.py index dd2f8af6..cf77c8df 100644 --- a/fbpcp/service/key_management_aws.py +++ b/fbpcp/service/key_management_aws.py @@ -9,16 +9,19 @@ from typing import Any, Dict, List, Optional from fbpcp.gateway.kms import KMSGateway +from fbpcp.service.key_management import KeyManagementService -class AWSKeyManagementService: +class AWSKeyManagementService(KeyManagementService): key_id: str + signing_algorithm: str grant_tokens: List[str] def __init__( self, region: str, key_id: str, + signing_algorithm: Optional[str] = None, grant_tokens: Optional[List[str]] = None, access_key_id: Optional[str] = None, access_key_data: Optional[str] = None, @@ -30,4 +33,17 @@ def __init__( """ self.kms_gateway = KMSGateway(region, access_key_id, access_key_data, config) self.key_id = key_id + self.signing_algorithm = signing_algorithm if signing_algorithm else "" self.grant_tokens = grant_tokens if grant_tokens else [] + + def sign(self, message: str, message_type: str = "RAW") -> str: + if not self.signing_algorithm: + raise ValueError("No Signing Algorithm Set") + signature = self.kms_gateway.sign( + key_id=self.key_id, + message=message, + message_type=message_type, + grant_tokens=self.grant_tokens, + signing_algorithm=self.signing_algorithm, + ) + return signature diff --git a/tests/gateway/test_kms.py b/tests/gateway/test_kms.py new file mode 100644 index 00000000..45476b54 --- /dev/null +++ b/tests/gateway/test_kms.py @@ -0,0 +1,47 @@ +#!/usr/bin/env fbpython +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import base64 as b64 +import unittest +from unittest.mock import MagicMock, patch + +from fbpcp.gateway.kms import KMSGateway + + +class TestKMSGateway(unittest.TestCase): + REGION = "us-west-2" + TEST_ACCESS_KEY_ID = "test-access-key-id" + TEST_ACCESS_KEY_DATA = "test-access-key-data" + + @patch("boto3.client") + def setUp(self, BotoClient) -> None: + self.kms = KMSGateway( + region=self.REGION, + access_key_id=self.TEST_ACCESS_KEY_ID, + access_key_data=self.TEST_ACCESS_KEY_DATA, + ) + self.kms.client = BotoClient() + + def test_sign(self) -> None: + # Arrange + sign_args = { + "key_id": "test_key_id", + "message": "test_message", + "message_type": "test_message_type", + "grant_tokens": [], + "signing_algorithm": "", + } + signed_message = "test_signed_message" + self.kms.client.sign = MagicMock( + return_value={"Signature": signed_message.encode()} + ) + + # Act + b64_signature = self.kms.sign(**sign_args) + signature = b64.b64decode(b64_signature.encode()).decode() + + # Assert + self.assertEqual(signature, signed_message) diff --git a/tests/service/test_key_managment_aws.py b/tests/service/test_key_managment_aws.py new file mode 100644 index 00000000..bb829412 --- /dev/null +++ b/tests/service/test_key_managment_aws.py @@ -0,0 +1,47 @@ +#!/usr/bin/env fbpython +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import unittest +from unittest.mock import MagicMock, patch + +from fbpcp.service.key_management_aws import AWSKeyManagementService + + +class TestAWSKeyManagementService(unittest.TestCase): + REGION = "us-west-2" + + TEST_KEY_ID = "test-key-id" + TEST_SIGNING_ALGORITHM = "test-signing-algorithm" + + TEST_ACCESS_KEY_ID = "test-access-key-id" + TEST_ACCESS_KEY_DATA = "test-access-key-data" + + @patch("boto3.client") + def setUp(self, BotoClient) -> None: + self.kms_aws = AWSKeyManagementService( + region=self.REGION, + key_id=self.TEST_KEY_ID, + signing_algorithm=self.TEST_SIGNING_ALGORITHM, + access_key_id=self.TEST_ACCESS_KEY_ID, + access_key_data=self.TEST_ACCESS_KEY_DATA, + ) + self.kms_aws.kms_gateway.client = BotoClient() + + def test_sign(self) -> None: + # Arrange + sign_args = { + "message": "test_message", + "message_type": "test_message_type", + } + signed_message = "test_signed_message" + + self.kms_aws.kms_gateway.sign = MagicMock(return_value=signed_message) + + # Act + signature = self.kms_aws.sign(**sign_args) + + # Assert + self.assertEqual(signature, signed_message) From e943ea612260b11036a719e854c5642d80429c8a Mon Sep 17 00:00:00 2001 From: Arnav Sacheti Date: Tue, 26 Jul 2022 13:59:08 -0700 Subject: [PATCH 3/4] Adding support for message verification (#378) Summary: Pull Request resolved: https://github.com/facebookresearch/fbpcp/pull/378 Differential Revision: D37667144 fbshipit-source-id: 32142cb73fe6d0f0a62090a18c3e455a5a4dc355 --- fbpcp/gateway/kms.py | 27 ++++++++++++++++++++++--- fbpcp/service/key_management.py | 4 ++++ fbpcp/service/key_management_aws.py | 11 ++++++++++ tests/gateway/test_kms.py | 19 +++++++++++++++++ tests/service/test_key_managment_aws.py | 22 +++++++++++++++++--- 5 files changed, 77 insertions(+), 6 deletions(-) diff --git a/fbpcp/gateway/kms.py b/fbpcp/gateway/kms.py index 40aefcb3..2f6efe88 100644 --- a/fbpcp/gateway/kms.py +++ b/fbpcp/gateway/kms.py @@ -6,7 +6,7 @@ # pyre-strict -from base64 import b64encode +from base64 import b64decode, b64encode from typing import Any, Dict, List, Optional import boto3 @@ -32,17 +32,38 @@ def __init__( def sign( self, key_id: str, - message: bytes, + message: str, message_type: str, grant_tokens: List[str], signing_algorithm: str, ) -> str: response = self.client.sign( KeyId=key_id, - Message=message, + Message=message.encode(), MessageType=message_type, GrantTokens=grant_tokens, SigningAlgorithm=signing_algorithm, ) signature = b64encode(response["Signature"]).decode() return signature + + @error_handler + def verify( + self, + key_id: str, + message: str, + message_type: str, + signature: str, + signing_algorithm: str, + grant_tokens: List[str], + ) -> bool: + b64_signature = b64decode(signature.encode()) + response = self.client.verify( + KeyId=key_id, + Message=message.encode(), + MessageType=message_type, + Signature=b64_signature, + SigningAlgorithm=signing_algorithm, + GrantTokens=grant_tokens, + ) + return response["SignatureValid"] diff --git a/fbpcp/service/key_management.py b/fbpcp/service/key_management.py index e9a83235..e8b50ada 100644 --- a/fbpcp/service/key_management.py +++ b/fbpcp/service/key_management.py @@ -13,3 +13,7 @@ class KeyManagementService(abc.ABC): @abc.abstractmethod def sign(self, message: str) -> str: pass + + @abc.abstractmethod + def verify(self, message: str, signature: str) -> bool: + pass diff --git a/fbpcp/service/key_management_aws.py b/fbpcp/service/key_management_aws.py index cf77c8df..75a58421 100644 --- a/fbpcp/service/key_management_aws.py +++ b/fbpcp/service/key_management_aws.py @@ -47,3 +47,14 @@ def sign(self, message: str, message_type: str = "RAW") -> str: signing_algorithm=self.signing_algorithm, ) return signature + + def verify(self, message: str, signature: str, message_type: str = "RAW") -> bool: + valid = self.kms_gateway.verify( + key_id=self.key_id, + message=message, + message_type=message_type, + signature=signature, + signing_algorithm=self.signing_algorithm, + grant_tokens=self.grant_tokens, + ) + return valid diff --git a/tests/gateway/test_kms.py b/tests/gateway/test_kms.py index 45476b54..c7e299c4 100644 --- a/tests/gateway/test_kms.py +++ b/tests/gateway/test_kms.py @@ -4,6 +4,7 @@ # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. +import base64 import base64 as b64 import unittest from unittest.mock import MagicMock, patch @@ -45,3 +46,21 @@ def test_sign(self) -> None: # Assert self.assertEqual(signature, signed_message) + + def test_verify(self) -> None: + # Arrange + verify_args = { + "key_id": "test_key_id", + "message": "test_message", + "message_type": "test_message_type", + "signature": "dGVzdF9tZXNzYWdl", + "grant_tokens": [], + "signing_algorithm": "", + } + self.kms.client.verify = MagicMock(return_value={"SignatureValid": True}) + + # Act + verification = self.kms.verify(**verify_args) + + # Assert + self.assertTrue(verification) diff --git a/tests/service/test_key_managment_aws.py b/tests/service/test_key_managment_aws.py index bb829412..55e39948 100644 --- a/tests/service/test_key_managment_aws.py +++ b/tests/service/test_key_managment_aws.py @@ -36,12 +36,28 @@ def test_sign(self) -> None: "message": "test_message", "message_type": "test_message_type", } - signed_message = "test_signed_message" + test_signature = "test_signature" - self.kms_aws.kms_gateway.sign = MagicMock(return_value=signed_message) + self.kms_aws.kms_gateway.sign = MagicMock(return_value=test_signature) # Act signature = self.kms_aws.sign(**sign_args) # Assert - self.assertEqual(signature, signed_message) + self.assertEqual(signature, test_signature) + + def test_verify(self) -> None: + # Arrange + verify_args = { + "signature": "test_signature", + "message": "test_message", + "message_type": "test_message_type", + } + + self.kms_aws.kms_gateway.verify = MagicMock(return_value=True) + + # Act + status = self.kms_aws.verify(**verify_args) + + # Assert + self.assertTrue(status) From 6862db8d369845024588f148dc951a3961f38853 Mon Sep 17 00:00:00 2001 From: Arnav Sacheti Date: Tue, 26 Jul 2022 14:00:33 -0700 Subject: [PATCH 4/4] Reorganizing objects created in CLI (#376) Summary: Pull Request resolved: https://github.com/facebookresearch/fbpcp/pull/376 Moving some imports to be inside of if statements when possible to move closer to only creating instances when needed Differential Revision: D38009745 fbshipit-source-id: 06fa60169209774dddeb7bd0591b6f644122a04d --- onedocker/script/cli/onedocker_cli.py | 68 ++++++++++++++++----------- 1 file changed, 40 insertions(+), 28 deletions(-) diff --git a/onedocker/script/cli/onedocker_cli.py b/onedocker/script/cli/onedocker_cli.py index 957b975f..4d2ef6b7 100644 --- a/onedocker/script/cli/onedocker_cli.py +++ b/onedocker/script/cli/onedocker_cli.py @@ -41,14 +41,14 @@ from onedocker.service.attestation import AttestationService logger = None -onedocker_svc = None -container_svc = None -onedocker_package_repo = None -onedocker_checksum_repo = None + attestation_svc = None +container_svc = None +onedocker_svc = None log_svc = None -task_definition = None -repository_path = None + +onedocker_checksum_repo = None +onedocker_package_repo = None DEFAULT_BINARY_VERSION = "latest" DEFAULT_TIMEOUT = 18000 @@ -171,12 +171,8 @@ def _build_log_service(config: Dict[str, Any]) -> LogService: return log_class(**config["constructor"]) -def _build_exe_s3_path(repository_path: str, package_name: str, version: str) -> str: - return f"{repository_path}{package_name}/{version}/{package_name.split('/')[-1]}" - - def main() -> None: - global container_svc, onedocker_svc, onedocker_package_repo, onedocker_checksum_repo, log_svc, logger, task_definition, repository_path, attestation_svc + global container_svc, onedocker_svc, onedocker_package_repo, onedocker_checksum_repo, log_svc, logger, attestation_svc s = schema.Schema( { "upload": bool, @@ -208,30 +204,46 @@ def main() -> None: version = ( arguments["--version"] if arguments["--version"] else DEFAULT_BINARY_VERSION ) - enable_attestation = arguments["--enable_attestation"] config = yaml.load(Path(arguments["--config"])).get("onedocker-cli") - task_definition = config["setting"]["task_definition"] - repository_path = config["setting"]["repository_path"] - checksum_repository_path = config["setting"].get("checksum_repository_path", "") - - attestation_svc = AttestationService() - storage_svc = _build_storage_service(config["dependency"]["StorageService"]) - container_svc = _build_container_service(config["dependency"]["ContainerService"]) - onedocker_svc = OneDockerService(container_svc, task_definition) - onedocker_package_repo = OneDockerPackageRepository(storage_svc, repository_path) - onedocker_checksum_repo = OneDockerChecksumRepository( - storage_svc, checksum_repository_path - ) - log_svc = _build_log_service(config["dependency"]["LogService"]) - status = "enabled" if enable_attestation else "disabled" - logger.info(f"Package tracking for package {package_name}: {version} is {status}") + if arguments["upload"] or arguments["show"]: + repository_path = config["setting"]["repository_path"] + checksum_repository_path = config["setting"].get("checksum_repository_path", "") + + storage_svc = _build_storage_service(config["dependency"]["StorageService"]) + + onedocker_package_repo = OneDockerPackageRepository( + storage_svc, repository_path + ) + if checksum_repository_path: + onedocker_checksum_repo = OneDockerChecksumRepository( + storage_svc, checksum_repository_path + ) + + if arguments["test"] or arguments["stop"]: + task_definition = config["setting"]["task_definition"] + container_svc = _build_container_service( + config["dependency"]["ContainerService"] + ) + onedocker_svc = OneDockerService(container_svc, task_definition) if arguments["upload"]: + enable_attestation = arguments["--enable_attestation"] + + status = "enabled" if enable_attestation else "disabled" + logger.info( + f"Package tracking for package {package_name}: {version} is {status}" + ) + + attestation_svc = AttestationService() + _upload(package_dir, package_name, version, enable_attestation) elif arguments["test"]: - timeout = arguments["--timeout"] if arguments["--timeout"] else DEFAULT_TIMEOUT + timeout = arguments.get("--timeout", DEFAULT_TIMEOUT) + + log_svc = _build_log_service(config["dependency"]["LogService"]) + _test(package_name, version, arguments["--cmd_args"], timeout) elif arguments["show"]: _show(package_name, arguments["--version"])