From 0b6ac2049390f942821a54900c4d5865998cba97 Mon Sep 17 00:00:00 2001 From: Arnav Sacheti Date: Thu, 14 Jul 2022 11:52:33 -0700 Subject: [PATCH 01/13] Adding function to handle uploading of checksums (#359) Summary: Pull Request resolved: https://github.com/facebookresearch/fbpcp/pull/359 # Context Want to move s3 i/o out of attestation service. # This commit Creates tooling that would enable move to happen, and changes checksum path to follow styling of repository path for simplicity Differential Revision: D37524841 fbshipit-source-id: ec1b84be5456c8340b8c3d8b1747dfc1e24e7085 --- onedocker/repository/onedocker_checksum.py | 30 +++++++++++ .../repository/test_onedocker_checksum.py | 54 +++++++++++++++++++ 2 files changed, 84 insertions(+) create mode 100644 onedocker/repository/onedocker_checksum.py create mode 100644 onedocker/tests/repository/test_onedocker_checksum.py diff --git a/onedocker/repository/onedocker_checksum.py b/onedocker/repository/onedocker_checksum.py new file mode 100644 index 00000000..67989693 --- /dev/null +++ b/onedocker/repository/onedocker_checksum.py @@ -0,0 +1,30 @@ +#!/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 fbpcp.service.storage import StorageService + + +class OneDockerChecksumRepository: + def __init__( + self, storage_svc: StorageService, checksum_repository_path: str + ) -> None: + self.storage_svc = storage_svc + self.checksum_repository_path = checksum_repository_path + + def _build_checksum_path(self, package_name: str, version: str) -> str: + if not self.checksum_repository_path: + raise ValueError( + "Checksum Repository Path not set. Unable to attest Package" + ) + return f"{self.checksum_repository_path}{package_name}/{version}/{package_name.split('/')[-1]}.json" + + def write(self, package_name: str, version: str, checksum_data: str) -> None: + package_path = self._build_checksum_path( + package_name=package_name, version=version + ) + self.storage_svc.write(filename=package_path, data=checksum_data) diff --git a/onedocker/tests/repository/test_onedocker_checksum.py b/onedocker/tests/repository/test_onedocker_checksum.py new file mode 100644 index 00000000..13cd217e --- /dev/null +++ b/onedocker/tests/repository/test_onedocker_checksum.py @@ -0,0 +1,54 @@ +#!/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. + + +import unittest +from unittest.mock import MagicMock, patch + +from onedocker.repository.onedocker_checksum import OneDockerChecksumRepository + + +class TestOneDockerChecksumRepository(unittest.TestCase): + TEST_PACKAGE_NAME = "project/exe_name" + TEST_PACKAGE_VERSION = "1.0" + + @patch("fbpcp.service.storage_s3.S3StorageService") + def setUp(self, MockStorageService): + self.checksum_repository_path = "/abc/" + self.onedocker_checksum = OneDockerChecksumRepository( + MockStorageService, self.checksum_repository_path + ) + self.expected_s3_dest = f"{self.checksum_repository_path}{self.TEST_PACKAGE_NAME}/{self.TEST_PACKAGE_VERSION}/{self.TEST_PACKAGE_NAME.split('/')[-1]}.json" + + def test_onedockerrepo_write(self): + # Arrange + checksum_data = "xyz" + + # Act + self.onedocker_checksum.write( + package_name=self.TEST_PACKAGE_NAME, + version=self.TEST_PACKAGE_VERSION, + checksum_data=checksum_data, + ) + + # Assert + self.onedocker_checksum.storage_svc.write.assert_called_with( + filename=self.expected_s3_dest, data=checksum_data + ) + + def test_onedockerrepo_write_no_checksum_path(self): + # Arrange + checksum_data = "xyz" + + onedocker_checksum = OneDockerChecksumRepository(MagicMock(), "") + + # Act & Assert + with self.assertRaises(ValueError): + onedocker_checksum.write( + package_name=self.TEST_PACKAGE_NAME, + version=self.TEST_PACKAGE_VERSION, + checksum_data=checksum_data, + ) From daae9d3df59e944a02f3d63f6d4bfa6240054852 Mon Sep 17 00:00:00 2001 From: Arnav Sacheti Date: Thu, 14 Jul 2022 11:52:33 -0700 Subject: [PATCH 02/13] Adding function to handle downloading of checksums Summary: Added functionalilty to upload in last commit, this extends this to also include downloading of checksum file Differential Revision: D37525237 fbshipit-source-id: 45ee1e3d3ad437f9aea39f6b5f6612f91428c76d --- onedocker/repository/onedocker_checksum.py | 18 +++++++++ .../repository/test_onedocker_checksum.py | 40 +++++++++++++++++++ 2 files changed, 58 insertions(+) diff --git a/onedocker/repository/onedocker_checksum.py b/onedocker/repository/onedocker_checksum.py index 67989693..04c9612d 100644 --- a/onedocker/repository/onedocker_checksum.py +++ b/onedocker/repository/onedocker_checksum.py @@ -23,8 +23,26 @@ def _build_checksum_path(self, package_name: str, version: str) -> str: ) return f"{self.checksum_repository_path}{package_name}/{version}/{package_name.split('/')[-1]}.json" + def _file_exists(self, package_name: str, version: str) -> bool: + package_path = self._build_checksum_path( + package_name=package_name, version=version + ) + return self.storage_svc.file_exists(package_path) + def write(self, package_name: str, version: str, checksum_data: str) -> None: package_path = self._build_checksum_path( package_name=package_name, version=version ) self.storage_svc.write(filename=package_path, data=checksum_data) + + def read(self, package_name: str, version: str) -> str: + package_path = self._build_checksum_path( + package_name=package_name, version=version + ) + + if not self._file_exists(package_name, version): + raise FileNotFoundError( + f"Cant find checksum file for package {package_name}, version {version}" + ) + + return self.storage_svc.read(filename=package_path) diff --git a/onedocker/tests/repository/test_onedocker_checksum.py b/onedocker/tests/repository/test_onedocker_checksum.py index 13cd217e..e1eed2b5 100644 --- a/onedocker/tests/repository/test_onedocker_checksum.py +++ b/onedocker/tests/repository/test_onedocker_checksum.py @@ -52,3 +52,43 @@ def test_onedockerrepo_write_no_checksum_path(self): version=self.TEST_PACKAGE_VERSION, checksum_data=checksum_data, ) + + def test_onedockerrepo_read(self): + # Arrange + checksum_data = "xyz" + + self.onedocker_checksum.storage_svc.read = MagicMock(return_value=checksum_data) + + # Act + actual_checksum_data = self.onedocker_checksum.read( + package_name=self.TEST_PACKAGE_NAME, + version=self.TEST_PACKAGE_VERSION, + ) + + # Assert + self.onedocker_checksum.storage_svc.read.assert_called_with( + filename=self.expected_s3_dest + ) + self.assertEqual(checksum_data, actual_checksum_data) + + def test_onedockerrepo_read_no_checksum_path(self): + # Arrange + onedocker_checksum = OneDockerChecksumRepository(MagicMock(), "") + + # Act & Assert + with self.assertRaises(ValueError): + onedocker_checksum.read( + package_name=self.TEST_PACKAGE_NAME, + version=self.TEST_PACKAGE_VERSION, + ) + + def test_onedockerrepo_read_no_file_exist(self): + # Arrange + self.onedocker_checksum.storage_svc.file_exists = MagicMock(return_value=False) + + # Act & Assert + with self.assertRaises(FileNotFoundError): + self.onedocker_checksum.read( + package_name=self.TEST_PACKAGE_NAME, + version=self.TEST_PACKAGE_VERSION, + ) From 8edbb212e8214d262dbf5ff7d5d94cb73a92b34e Mon Sep 17 00:00:00 2001 From: Arnav Sacheti Date: Thu, 14 Jul 2022 11:52:33 -0700 Subject: [PATCH 03/13] Moving checksum upload from attestation to ODCR (#360) Summary: Pull Request resolved: https://github.com/facebookresearch/fbpcp/pull/360 # Context We added functionality to upload checksums to ODPR, but dont use it atm. # This commit Removes the uploading from AttestationService and puts it into ODPR. Differential Revision: D37526228 fbshipit-source-id: 3f12457300792b9c5a93e182efbcf7a40054722a --- onedocker/script/cli/onedocker_cli.py | 33 ++++++++++----- onedocker/service/attestation.py | 40 +++++++------------ .../tests/script/cli/test_onedocker_cli.py | 18 ++++++++- onedocker/tests/service/test_attestation.py | 15 ++----- 4 files changed, 57 insertions(+), 49 deletions(-) diff --git a/onedocker/script/cli/onedocker_cli.py b/onedocker/script/cli/onedocker_cli.py index 6063e335..bfdc83ba 100644 --- a/onedocker/script/cli/onedocker_cli.py +++ b/onedocker/script/cli/onedocker_cli.py @@ -36,6 +36,7 @@ from fbpcp.service.onedocker import OneDockerService from fbpcp.service.storage import StorageService from fbpcp.util import reflect, yaml +from onedocker.repository.onedocker_checksum import OneDockerChecksumRepository from onedocker.repository.onedocker_package import OneDockerPackageRepository from onedocker.service.attestation import AttestationService @@ -43,7 +44,8 @@ onedocker_svc = None container_svc = None onedocker_package_repo = None -attestation_service = None +onedocker_checksum_repo = None +attestation_svc = None log_svc = None task_definition = None repository_path = None @@ -65,7 +67,16 @@ def _upload( logger.info( f"Generating and uploading checksums for package {package_name}: {version}" ) - attestation_service.track_binary(package_dir, package_name, version) + formated_checksum_info = attestation_svc.track_binary( + binary_path=package_dir, + package_name=package_name, + version=version, + ) + onedocker_checksum_repo.write( + package_name=package_name, + version=version, + checksum_data=formated_checksum_info, + ) onedocker_package_repo.upload(package_name, version, package_dir) logger.info(f" Finished uploading '{package_name}, version {version}'.\n") @@ -165,7 +176,7 @@ def _build_exe_s3_path(repository_path: str, package_name: str, version: str) -> def main() -> None: - global container_svc, onedocker_svc, onedocker_package_repo, log_svc, logger, task_definition, repository_path, attestation_service + global container_svc, onedocker_svc, onedocker_package_repo, onedocker_checksum_repo, log_svc, logger, task_definition, repository_path, attestation_svc s = schema.Schema( { "upload": bool, @@ -197,24 +208,26 @@ 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"]) - enable_attestation = arguments["--enable_attestation"] - if enable_attestation: - logger.info( - f"Package tracking for package {package_name}: {version} has been enabled" - ) - checksum_repository_path = config["setting"]["checksum_repository_path"] - attestation_service = AttestationService(storage_svc, checksum_repository_path) + logger.info( + f"Package tracking for package {package_name}: {version} has been {'enabled' if enable_attestation else 'disabled'}" + ) if arguments["upload"]: _upload(package_dir, package_name, version, enable_attestation) diff --git a/onedocker/service/attestation.py b/onedocker/service/attestation.py index 82bb160d..7f4ac9a3 100644 --- a/onedocker/service/attestation.py +++ b/onedocker/service/attestation.py @@ -8,7 +8,7 @@ import json import logging -from typing import Any, Dict, List +from typing import Dict, List, Optional from fbpcp.service.storage import StorageService from onedocker.entity.attestation_error import AttestationError @@ -25,11 +25,16 @@ class AttestationService: ChecksumType.BLAKE2B, ] - def __init__(self, storage_svc: StorageService, repository_path: str) -> None: + def __init__( + self, + storage_svc: Optional[StorageService] = None, + repository_path: str = "", + ) -> None: self.logger: logging.Logger = logging.getLogger(__name__) self.checksum_generator = LocalChecksumGenerator() - self.storage_svc = storage_svc - self.repository_path = repository_path + if storage_svc is not None: + self.storage_svc: StorageService = storage_svc + self.repository_path: str = repository_path def _build_attestation_repository_path( self, @@ -56,26 +61,12 @@ def _get_checksum_info( checksums=checksums, ) - def _upload_checksum( - self, - package_name: str, - version: str, - checksum_info: Dict[str, Any], - ) -> None: - - # Construct file information - (name and contents) - file_path = self._build_attestation_repository_path(package_name, version) - file_contents = json.dumps(checksum_info, indent=4) - - # upload contents to set repo path - self.storage_svc.write(file_path, file_contents) - def track_binary( self, binary_path: str, package_name: str, version: str, - ) -> None: + ) -> str: """ This Function generates then uploads checksums for passed in local binary @@ -83,6 +74,9 @@ def track_binary( binary_path: Local file path pointing to the package package_name: Package Name to use when uploading file to checksum repository version: Package Version to relay while uploading file to checksum repository + + Returns: + formated_checksum_info: A JSON formated file that contains all the checksum data for a file """ # Generates checksums self.logger.info(f"Generating checksums for binary at {binary_path}") @@ -92,13 +86,7 @@ def track_binary( binary_path=binary_path, ) - # Upload checksums and package info to set repo path - self.logger.info(f"Uploading checksums for package {package_name}: {version}") - self._upload_checksum( - package_name=package_name, - version=version, - checksum_info=checksum_info.asdict(), - ) + return json.dumps(checksum_info.asdict(), indent=4) def attest_binary( self, diff --git a/onedocker/tests/script/cli/test_onedocker_cli.py b/onedocker/tests/script/cli/test_onedocker_cli.py index 00b04610..82d9ee3c 100644 --- a/onedocker/tests/script/cli/test_onedocker_cli.py +++ b/onedocker/tests/script/cli/test_onedocker_cli.py @@ -17,6 +17,7 @@ from fbpcp.service.onedocker import OneDockerService from fbpcp.util import yaml as util_yaml from onedocker.entity.package_info import PackageInfo +from onedocker.repository.onedocker_checksum import OneDockerChecksumRepository from onedocker.repository.onedocker_package import OneDockerPackageRepository from onedocker.script.cli.onedocker_cli import __doc__ as __onedocker_cli_doc__, main from onedocker.service.attestation import AttestationService @@ -57,6 +58,7 @@ def setUp(self): self.timeout = "100" self.cmd_args = "-h" self.container = "secret_container" + self.checksums = "formated_checksums_go_here" self.package_info = PackageInfo( package_name=self.package_name, @@ -101,6 +103,11 @@ def setUp(self): "upload", MagicMock(return_value=None), ).start() + self.mockODCRWrite = patch.object( + OneDockerChecksumRepository, + "write", + MagicMock(return_value=None), + ).start() self.mockODPRGetPackageVersions = patch.object( OneDockerPackageRepository, "get_package_versions", @@ -115,7 +122,7 @@ def setUp(self): self.mockAttestationServiceTrackBinary = patch.object( AttestationService, "track_binary", - MagicMock(), + MagicMock(return_value=self.checksums), ).start() self.mockContainerService = patch.object( @@ -286,11 +293,18 @@ def test_upload_attestation_service(self): # Assert self.mockYamlLoad.assert_called_once() + self.mockODCRWrite.assert_called_once_with( + package_name=self.package_name, + version=self.version, + checksum_data=self.checksums, + ) self.mockODPRUpload.assert_called_once_with( self.package_name, self.version, self.package_dir ) self.mockAttestationServiceTrackBinary.assert_called_once_with( - self.package_dir, self.package_name, self.version + binary_path=self.package_dir, + package_name=self.package_name, + version=self.version, ) @patch.object(CloudWatchLogService, "get_log_path") diff --git a/onedocker/tests/service/test_attestation.py b/onedocker/tests/service/test_attestation.py index f5f03b8a..6591dba6 100644 --- a/onedocker/tests/service/test_attestation.py +++ b/onedocker/tests/service/test_attestation.py @@ -6,7 +6,7 @@ import unittest from json import dumps -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock from fbpcp.service.storage_s3 import S3StorageService from onedocker.entity.attestation_error import AttestationError @@ -55,13 +55,9 @@ def setUp(self) -> None: ) self.attestation_service.storage_svc.file_exists = MagicMock(return_value=True) - @patch.object(S3StorageService, "write") - def test_track_binary_s3( - self, - mockS3StorageServiceWrite, - ): + def test_track_binary_s3(self): # Arrange & Act - self.attestation_service.track_binary( + formated_checksums = self.attestation_service.track_binary( binary_path=self.test_package["binary_path"], package_name=self.test_package["name"], version=self.test_package["version"], @@ -72,10 +68,7 @@ def test_track_binary_s3( binary_path=self.test_package["binary_path"], checksum_algorithms=self.algorithms, ) - mockS3StorageServiceWrite.assert_called_once_with( - self.test_package["checksum_path"], - self.file_contents, - ) + self.assertEqual(formated_checksums, self.file_contents) def test_attest_binary_s3( self, From 7bcbc6955b531ace4e951eb5e5c993449986585c Mon Sep 17 00:00:00 2001 From: Arnav Sacheti Date: Thu, 14 Jul 2022 11:52:33 -0700 Subject: [PATCH 04/13] Moving checksum download from attestation to ODCR Summary: # Context Want to move all s3 i/o out of AttestationService # This Commit Moves download from AttestationService to ODPR Differential Revision: D37526823 fbshipit-source-id: ce810ae53fce9af4b5bd5d364376f819b5524226 --- onedocker/script/runner/onedocker_runner.py | 9 +- onedocker/service/attestation.py | 41 ++------- onedocker/tests/service/test_attestation.py | 96 ++++++--------------- 3 files changed, 42 insertions(+), 104 deletions(-) diff --git a/onedocker/script/runner/onedocker_runner.py b/onedocker/script/runner/onedocker_runner.py index 25bc6f96..d7041f3d 100644 --- a/onedocker/script/runner/onedocker_runner.py +++ b/onedocker/script/runner/onedocker_runner.py @@ -48,6 +48,7 @@ ) from onedocker.common.util import run_cmd from onedocker.entity.checksum_type import ChecksumType +from onedocker.repository.onedocker_checksum import OneDockerChecksumRepository from onedocker.repository.onedocker_package import OneDockerPackageRepository from onedocker.service.attestation import AttestationService from onedocker.service.certificate_self_signed import SelfSignedCertificateService @@ -246,12 +247,18 @@ def _attest_executable( f"Starting verification for package {package_name}: {version} using checksum type: {checksum_type.name}" ) storage_svc = S3StorageService(S3Path(checksum_repository_path).region) - attestation_service = AttestationService(storage_svc, checksum_repository_path) + attestation_service = AttestationService() + onedocker_checksum_repository = OneDockerChecksumRepository( + storage_svc, checksum_repository_path + ) + + formated_checksum_info = onedocker_checksum_repository.read(package_name, version) attestation_service.attest_binary( binary_path=binary_path, package_name=package_name, version=version, + formated_checksum_info=formated_checksum_info, checksum_algorithm=checksum_type, ) logger.info(f"Finished verification for package {package_name}: {version}") diff --git a/onedocker/service/attestation.py b/onedocker/service/attestation.py index 7f4ac9a3..11ebbf71 100644 --- a/onedocker/service/attestation.py +++ b/onedocker/service/attestation.py @@ -8,9 +8,8 @@ import json import logging -from typing import Dict, List, Optional +from typing import Dict, List -from fbpcp.service.storage import StorageService from onedocker.entity.attestation_error import AttestationError from onedocker.entity.checksum_info import ChecksumInfo from onedocker.entity.checksum_type import ChecksumType @@ -25,23 +24,9 @@ class AttestationService: ChecksumType.BLAKE2B, ] - def __init__( - self, - storage_svc: Optional[StorageService] = None, - repository_path: str = "", - ) -> None: + def __init__(self) -> None: self.logger: logging.Logger = logging.getLogger(__name__) self.checksum_generator = LocalChecksumGenerator() - if storage_svc is not None: - self.storage_svc: StorageService = storage_svc - self.repository_path: str = repository_path - - def _build_attestation_repository_path( - self, - package_name: str, - version: str, - ) -> str: - return f"{self.repository_path}{package_name}/{version}.json" def _get_checksum_info( self, @@ -93,6 +78,7 @@ def attest_binary( binary_path: str, package_name: str, version: str, + formated_checksum_info: str, checksum_algorithm: ChecksumType, ) -> None: """ @@ -101,32 +87,23 @@ def attest_binary( Args: binary_path: Local file path pointing to the package package_name: Package Name to use when downlading the checksum file from checksum repository - version: Package Version to relay while downloading the checksum file from checksum repository + version: Package Version to relay while downloading the checksum file from checksum repository + formated_checksum_info: String encoding of ChecksumInfo attaining to the JSON file format checksum_algorithm: Checksum algorithm that should be used while attesting local binary """ - # Build file path - file_path = self._build_attestation_repository_path(package_name, version) - - if not self.storage_svc.file_exists(file_path): - self.logger.info( - f"Untracked package {package_name}: {version}. Skipping Attestion." - ) - return None - # Retrieve file info and parse contents - file_contents = self.storage_svc.read(file_path) - package_info = json.loads(file_contents) - package_checksum_info = ChecksumInfo(**package_info) + checksum_info_dict = json.loads(formated_checksum_info) + checksum_info = ChecksumInfo(**checksum_info_dict) # Generates checksums self.logger.info(f"Generating checksums for binary at {binary_path}") - checksum_info = self._get_checksum_info( + binary_checksum_info = self._get_checksum_info( package_name=package_name, version=version, binary_path=binary_path, checksum_types=[checksum_algorithm], ) - if checksum_info != package_checksum_info: + if binary_checksum_info != checksum_info: raise AttestationError( "Downloaded binaries checksum information differs from uploaded package's checksum information" ) diff --git a/onedocker/tests/service/test_attestation.py b/onedocker/tests/service/test_attestation.py index 6591dba6..0b64802e 100644 --- a/onedocker/tests/service/test_attestation.py +++ b/onedocker/tests/service/test_attestation.py @@ -8,7 +8,6 @@ from json import dumps from unittest.mock import MagicMock -from fbpcp.service.storage_s3 import S3StorageService from onedocker.entity.attestation_error import AttestationError from onedocker.entity.checksum_info import ChecksumInfo from onedocker.entity.checksum_type import ChecksumType @@ -18,14 +17,9 @@ class TestAttestationService(unittest.TestCase): def setUp(self) -> None: # Globals varibales for tests - self.repository_path = ( - "https://onedocker-runner-unittest-asacheti.s3.us-west-2.amazonaws.com/" - ) - checksum_path = f"{self.repository_path}ls/latest.json" self.test_package = { "binary_path": "/usr/bin/ls", - "checksum_path": checksum_path, - "name": "ls", + "package_name": "ls", "version": "latest", } self.algorithms = list(ChecksumType) @@ -33,13 +27,10 @@ def setUp(self) -> None: k.name: f"valid_{k.name.lower()}_checksum_goes_here" for k in self.algorithms } - self.attestation_service = AttestationService( - S3StorageService("us-west-2"), - self.repository_path, - ) + self.attestation_service = AttestationService() self.file_contents = dumps( ChecksumInfo( - package_name=self.test_package["name"], + package_name=self.test_package["package_name"], version=self.test_package["version"], checksums=self.checksums, ).asdict(), @@ -50,16 +41,12 @@ def setUp(self) -> None: self.attestation_service.checksum_generator.generate_checksums = MagicMock( return_value=self.checksums ) - self.attestation_service.storage_svc.read = MagicMock( - return_value=self.file_contents - ) - self.attestation_service.storage_svc.file_exists = MagicMock(return_value=True) def test_track_binary_s3(self): # Arrange & Act formated_checksums = self.attestation_service.track_binary( binary_path=self.test_package["binary_path"], - package_name=self.test_package["name"], + package_name=self.test_package["package_name"], version=self.test_package["version"], ) @@ -84,8 +71,9 @@ def test_attest_binary_s3( # Act self.attestation_service.attest_binary( binary_path=self.test_package["binary_path"], - package_name=self.test_package["name"], + package_name=self.test_package["package_name"], version=self.test_package["version"], + formated_checksum_info=self.file_contents, checksum_algorithm=test_algorithm, ) @@ -94,12 +82,6 @@ def test_attest_binary_s3( binary_path=self.test_package["binary_path"], checksum_algorithms=[test_algorithm], ) - self.attestation_service.storage_svc.read.assert_called_once_with( - self.test_package["checksum_path"], - ) - self.attestation_service.storage_svc.file_exists.assert_called_once_with( - self.test_package["checksum_path"], - ) def test_attest_binary_s3_nonmatching_algorithm( self, @@ -108,11 +90,9 @@ def test_attest_binary_s3_nonmatching_algorithm( checksum_key = list(self.checksums.keys())[-1] # Should be blake2b test_algorithm = ChecksumType(checksum_key) - self.attestation_service.storage_svc.read.return_value = ( - self.file_contents.replace( - checksum_key.upper(), - "def_not_" + checksum_key.upper(), - ) + modified_file_contents = self.file_contents.replace( + checksum_key.upper(), + "def_not_" + checksum_key.upper(), ) # Modify file contents to be missing checksum causing failure self.attestation_service.checksum_generator.generate_checksums.return_value = { @@ -123,29 +103,22 @@ def test_attest_binary_s3_nonmatching_algorithm( with self.assertRaises(ValueError): self.attestation_service.attest_binary( binary_path=self.test_package["binary_path"], - package_name=self.test_package["name"], + package_name=self.test_package["package_name"], version=self.test_package["version"], + formated_checksum_info=modified_file_contents, checksum_algorithm=test_algorithm, ) # Assert assert not self.attestation_service.checksum_generator.generate_checksums.called - self.attestation_service.storage_svc.read.assert_called_once_with( - self.test_package["checksum_path"], - ) - self.attestation_service.storage_svc.file_exists.assert_called_once_with( - self.test_package["checksum_path"], - ) def test_attest_binary_s3_bad_name( self, ): # Arrange - self.attestation_service.storage_svc.read.return_value = ( - self.file_contents.replace( - self.test_package["name"], - f'def_not_{self.test_package["name"]}', - ) + modified_file_contents = self.file_contents.replace( + self.test_package["package_name"], + f'def_not_{self.test_package["package_name"]}', ) # Modifying package name to cause failure checksum_key = list(self.checksums.keys())[0] # Should be MD5 @@ -155,8 +128,9 @@ def test_attest_binary_s3_bad_name( with self.assertRaises(AttestationError): self.attestation_service.attest_binary( binary_path=self.test_package["binary_path"], - package_name=self.test_package["name"], + package_name=self.test_package["package_name"], version=self.test_package["version"], + formated_checksum_info=modified_file_contents, checksum_algorithm=test_algorithm, ) @@ -165,22 +139,14 @@ def test_attest_binary_s3_bad_name( binary_path=self.test_package["binary_path"], checksum_algorithms=[test_algorithm], ) - self.attestation_service.storage_svc.read.assert_called_once_with( - self.test_package["checksum_path"], - ) - self.attestation_service.storage_svc.file_exists.assert_called_once_with( - self.test_package["checksum_path"], - ) def test_attest_binary_s3_bad_version( self, ): # Arrange - self.attestation_service.storage_svc.read.return_value = ( - self.file_contents.replace( - self.test_package["version"], - f'def_not_{self.test_package["version"]}', - ) + modified_file_contents = self.file_contents.replace( + self.test_package["version"], + f'def_not_{self.test_package["version"]}', ) # Modifying package version to cause failure checksum_key = list(self.checksums.keys())[0] # Should be MD5 @@ -190,8 +156,9 @@ def test_attest_binary_s3_bad_version( with self.assertRaises(AttestationError): self.attestation_service.attest_binary( binary_path=self.test_package["binary_path"], - package_name=self.test_package["name"], + package_name=self.test_package["package_name"], version=self.test_package["version"], + formated_checksum_info=modified_file_contents, checksum_algorithm=test_algorithm, ) @@ -200,12 +167,6 @@ def test_attest_binary_s3_bad_version( binary_path=self.test_package["binary_path"], checksum_algorithms=[test_algorithm], ) - self.attestation_service.storage_svc.read.assert_called_once_with( - self.test_package["checksum_path"], - ) - self.attestation_service.storage_svc.file_exists.assert_called_once_with( - self.test_package["checksum_path"], - ) def test_attest_binary_s3_bad_checksum( self, @@ -214,10 +175,8 @@ def test_attest_binary_s3_bad_checksum( checksum_key = list(self.checksums.keys())[0] # Should be MD5 test_algorithm = ChecksumType(checksum_key) - self.attestation_service.storage_svc.read.return_value = ( - self.file_contents.replace( - "valid_" + checksum_key.lower(), "invalid_" + checksum_key.lower() - ) + modified_file_contents = self.file_contents.replace( + "valid_" + checksum_key.lower(), "invalid_" + checksum_key.lower() ) # Modifying md5 checksum to cause failure self.attestation_service.checksum_generator.generate_checksums.return_value = { @@ -228,8 +187,9 @@ def test_attest_binary_s3_bad_checksum( with self.assertRaises(AttestationError): self.attestation_service.attest_binary( binary_path=self.test_package["binary_path"], - package_name=self.test_package["name"], + package_name=self.test_package["package_name"], version=self.test_package["version"], + formated_checksum_info=modified_file_contents, checksum_algorithm=test_algorithm, ) @@ -238,9 +198,3 @@ def test_attest_binary_s3_bad_checksum( binary_path=self.test_package["binary_path"], checksum_algorithms=[test_algorithm], ) - self.attestation_service.storage_svc.read.assert_called_once_with( - self.test_package["checksum_path"], - ) - self.attestation_service.storage_svc.file_exists.assert_called_once_with( - self.test_package["checksum_path"], - ) From ce8a685f855c1f76969c544349091c4dc958e181 Mon Sep 17 00:00:00 2001 From: Arnav Sacheti Date: Thu, 14 Jul 2022 11:52:33 -0700 Subject: [PATCH 05/13] Addting test to keep test coverage above 80% Differential Revision: D37566562 fbshipit-source-id: 76271b94c30e6bd28a0091175c5c1b129a0ea8d0 --- .../script/runner/test_onedocker_runner.py | 45 +++++++++++++++++-- 1 file changed, 42 insertions(+), 3 deletions(-) diff --git a/onedocker/tests/script/runner/test_onedocker_runner.py b/onedocker/tests/script/runner/test_onedocker_runner.py index 89e4bac3..a1626479 100644 --- a/onedocker/tests/script/runner/test_onedocker_runner.py +++ b/onedocker/tests/script/runner/test_onedocker_runner.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 signal import sys import unittest from unittest.mock import patch @@ -11,6 +12,7 @@ from docopt import docopt from fbpcp.entity.certificate_request import CertificateRequest, KeyAlgorithm from fbpcp.error.pcp import InvalidParameterError +from onedocker.common.core_dump_handler_aws import AWSCoreDumpHandler from onedocker.repository.onedocker_package import OneDockerPackageRepository from onedocker.script.runner.onedocker_runner import ( __doc__ as __onedocker_runner_doc__, @@ -252,13 +254,50 @@ def test_main_env(self): # Assert self.assertEqual(cm.exception.code, 0) + @patch("onedocker.script.runner.onedocker_runner.run_cmd") + @patch.object(AWSCoreDumpHandler, "locate_core_dump_file") + @patch.object(AWSCoreDumpHandler, "upload_core_dump_file") + def test_main_core_dump( + self, mockUploadCoreDumpFile, mockLocateCoreDumpFile, mockRunCMD + ): + # Arrange & Act + error_code = 128 + signal.SIGSEGV + mockRunCMD.return_value = error_code + core_dump_file = "core_dump.txt" + mockLocateCoreDumpFile.return_value = core_dump_file + + with patch.object( + sys, + "argv", + [ + "onedocker-runner", + "echo", + "--version=latest", + '--exe_args="Test Message"', + "--exe_path=/usr/bin/", + "--repository_path=local", + ], + ): + with patch( + "os.getenv", + side_effect=lambda x: getenv(x), + ): + with self.assertRaises(SystemExit) as cm: + main() + # Assert + self.assertEqual(cm.exception.code, error_code) + mockLocateCoreDumpFile.assert_called_once() + mockUploadCoreDumpFile.assert_called_once_with( + core_dump_file, getenv("CORE_DUMP_REPOSITORY_PATH") + ) + def getenv(key): if key == "ONEDOCKER_REPOSITORY_PATH": - return "https://onedocker-runner-unittest-asacheti.s3.us-west-2.amazonaws.com/" + return "https://onedocker-package-repo.s3.us-west-2.amazonaws.com/" elif key == "ONEDOCKER_CHECKSUM_REPOSITORY_PATH": - return "https://onedocker-checksum-test.s3.us-west-2.amazonaws.com/checksums/" + return "https://onedocker-checksum-repo.s3.us-west-2.amazonaws.com/checksums/" elif key == "CORE_DUMP_REPOSITORY_PATH": - return "~/fbsource/fbcode/measurement/private_measurement/pcp/oss/onedocker/tests/script/runner/core_dump/" + return "https://onedocker-core-dump.s3.us-west-2.amazonaws.com/checksums/" else: return None From 8500de50f840cbb31c505b8d6acde50212dbef35 Mon Sep 17 00:00:00 2001 From: Arnav Sacheti Date: Thu, 14 Jul 2022 11:52:33 -0700 Subject: [PATCH 06/13] Cleaning up some logs Summary: Just adding some logs to help clean up output Differential Revision: D37527124 fbshipit-source-id: 20b074fcb8419ee416860cdbf33fa5ce24ca7339 --- onedocker/entity/checksum_info.py | 6 +++--- onedocker/script/cli/onedocker_cli.py | 11 +++++------ onedocker/script/runner/onedocker_runner.py | 6 ++++-- onedocker/tests/script/cli/test_onedocker_cli.py | 2 ++ 4 files changed, 14 insertions(+), 11 deletions(-) diff --git a/onedocker/entity/checksum_info.py b/onedocker/entity/checksum_info.py index 66ff933b..09458dad 100644 --- a/onedocker/entity/checksum_info.py +++ b/onedocker/entity/checksum_info.py @@ -20,8 +20,8 @@ class ChecksumInfo: This dataclass tracks a package's checksum info for attestation, to allow for easy comparision and tracking Fields: - package_name: String containting the package's name - version: String containting the package's version + package_name: String containing the package name + version: String containing the package version Checksums: Dict that holds a pairing between a ChecksumType (Key) and the corresponding hash (Value) """ @@ -41,7 +41,7 @@ def __post_init__(self) -> None: def __eq__(self, other: ChecksumInfo) -> bool: """ - Compares two ChecksumInfo isntannces in previously decided order + Compares two ChecksumInfo instances in previously decided order 1. Compares package_name 2. Compares version 3. Checks if overlaping checkum algorithms are present diff --git a/onedocker/script/cli/onedocker_cli.py b/onedocker/script/cli/onedocker_cli.py index bfdc83ba..54804b5f 100644 --- a/onedocker/script/cli/onedocker_cli.py +++ b/onedocker/script/cli/onedocker_cli.py @@ -64,19 +64,19 @@ def _upload( f" Starting uploading package {package_name} at '{package_dir}', version {version}..." ) if enable_attestation: - logger.info( - f"Generating and uploading checksums for package {package_name}: {version}" - ) + logger.info(f"Generating checksums for package {package_name}: {version}") formated_checksum_info = attestation_svc.track_binary( binary_path=package_dir, package_name=package_name, version=version, ) + logger.info(f"Uploading checksums for package {package_name}: {version}") onedocker_checksum_repo.write( package_name=package_name, version=version, checksum_data=formated_checksum_info, ) + logger.info(f"Uploading binary for package {package_name}: {version}") onedocker_package_repo.upload(package_name, version, package_dir) logger.info(f" Finished uploading '{package_name}, version {version}'.\n") @@ -225,9 +225,8 @@ def main() -> None: ) log_svc = _build_log_service(config["dependency"]["LogService"]) - logger.info( - f"Package tracking for package {package_name}: {version} has been {'enabled' if enable_attestation else 'disabled'}" - ) + status = "enabled" if enable_attestation else "disabled" + logger.info(f"Package tracking for package {package_name}: {version} is {status}") if arguments["upload"]: _upload(package_dir, package_name, version, enable_attestation) diff --git a/onedocker/script/runner/onedocker_runner.py b/onedocker/script/runner/onedocker_runner.py index d7041f3d..d7e1b722 100644 --- a/onedocker/script/runner/onedocker_runner.py +++ b/onedocker/script/runner/onedocker_runner.py @@ -244,7 +244,7 @@ def _attest_executable( version: str, ) -> None: logger.info( - f"Starting verification for package {package_name}: {version} using checksum type: {checksum_type.name}" + f"Starting attestation for package {package_name}: {version} using checksum type: {checksum_type.name}" ) storage_svc = S3StorageService(S3Path(checksum_repository_path).region) attestation_service = AttestationService() @@ -252,8 +252,10 @@ def _attest_executable( storage_svc, checksum_repository_path ) + logger.info(f"Downloading checksum info for package {package_name}: {version}") formated_checksum_info = onedocker_checksum_repository.read(package_name, version) + logger.info(f"Attesting checksum info for package {package_name}: {version}") attestation_service.attest_binary( binary_path=binary_path, package_name=package_name, @@ -261,7 +263,7 @@ def _attest_executable( formated_checksum_info=formated_checksum_info, checksum_algorithm=checksum_type, ) - logger.info(f"Finished verification for package {package_name}: {version}") + logger.info(f"Finished attestation for package {package_name}: {version}") def _parse_package_name(package_name: str) -> str: diff --git a/onedocker/tests/script/cli/test_onedocker_cli.py b/onedocker/tests/script/cli/test_onedocker_cli.py index 82d9ee3c..e0ec0548 100644 --- a/onedocker/tests/script/cli/test_onedocker_cli.py +++ b/onedocker/tests/script/cli/test_onedocker_cli.py @@ -379,6 +379,7 @@ def test_show(self): ): main() + # Assert self.mockYamlLoad.assert_called_once() self.mockODPRGetPackageInfo.assert_called_once_with( self.package_name, self.version @@ -398,6 +399,7 @@ def test_show_no_version(self): ): main() + # Assert self.mockYamlLoad.assert_called_once() self.mockODPRGetPackageVersions.assert_called_once_with(self.package_name) self.mockODPRGetPackageInfo.assert_called_once_with( From b3bf32f86f189ea99951624266350e78d7254ed4 Mon Sep 17 00:00:00 2001 From: Arnav Sacheti Date: Thu, 14 Jul 2022 11:52:33 -0700 Subject: [PATCH 07/13] Integrate into onedocker-runner cli (#344) Summary: Pull Request resolved: https://github.com/facebookresearch/fbpcp/pull/344 # Context What good is integration if no easy way to interact when running program from CLI # This commit Adds default and env variable lookup allowing run from CLI Differential Revision: D37229577 fbshipit-source-id: 625101ea81d22dce2423cd7811e360c8d9f125fd --- onedocker/common/env.py | 6 ++ onedocker/script/runner/onedocker_runner.py | 97 +++++++++++++------ .../script/runner/test_onedocker_runner.py | 11 +-- 3 files changed, 76 insertions(+), 38 deletions(-) diff --git a/onedocker/common/env.py b/onedocker/common/env.py index d34e20bf..23985ae8 100644 --- a/onedocker/common/env.py +++ b/onedocker/common/env.py @@ -7,8 +7,14 @@ # This is the repository path that OneDocker downloads binaries from ONEDOCKER_REPOSITORY_PATH = "ONEDOCKER_REPOSITORY_PATH" +# This is the repository path that OneDocker downloads binary checksums from +ONEDOCKER_CHECKSUM_REPOSITORY_PATH = "ONEDOCKER_CHECKSUM_REPOSITORY_PATH" + # This is the local path that the binaries reside ONEDOCKER_EXE_PATH = "ONEDOCKER_EXE_PATH" # This is the path user specified to upload the core dump file to CORE_DUMP_REPOSITORY_PATH = "CORE_DUMP_REPOSITORY_PATH" + +# This is the type of checksum we want to compare when running program +ONEDOCKER_CHECKSUM_TYPE = "ONEDOCKER_CHECKSUM_TYPE" diff --git a/onedocker/script/runner/onedocker_runner.py b/onedocker/script/runner/onedocker_runner.py index d7e1b722..da556d88 100644 --- a/onedocker/script/runner/onedocker_runner.py +++ b/onedocker/script/runner/onedocker_runner.py @@ -14,14 +14,16 @@ onedocker-runner --version= [options] Options: - -h --help Show this help - --repository_path= OneDocker repository path where the executables are downloaded from. No download when "LOCAL" repository is specified. - --exe_path= The local path where the executables are downloaded to. - --exe_args= The arguments the executable will use. - --timeout= Set timeout (in sec) to kill the task. - --log_path= Override the default path where logs are saved. - --cert_params= string format of CertificateRequest dictionary if a TLS certificate is requested - --verbose Set logging level to DEBUG. + -h --help Show this help + --repository_path= OneDocker repository path where the executables are downloaded from. No download when "LOCAL" repository is specified. + --checksum_repository_path= OneDocker checksum path where the checksums are downloaded from. Doesnt verify files when not specified + --exe_path= The local path where the executables are downloaded to. + --exe_args= The arguments the executable will use. + --checksum_type= Type of Checksum to use while attesting binary (Supported Types are: MD5, SHA256[Default], BLAKE2B) + --timeout= Set timeout (in sec) to kill the task. + --log_path= Override the default path where logs are saved. + --cert_params= String format of CertificateRequest dictionary if a TLS certificate is requested + --verbose Set logging level to DEBUG. """ import logging import os @@ -38,11 +40,15 @@ import schema from docopt import docopt from fbpcp.entity.certificate_request import CertificateRequest + +from fbpcp.error.pcp import PcpError from fbpcp.service.storage_s3 import S3StorageService from fbpcp.util.s3path import S3Path from onedocker.common.core_dump_handler_aws import AWSCoreDumpHandler from onedocker.common.env import ( CORE_DUMP_REPOSITORY_PATH, + ONEDOCKER_CHECKSUM_REPOSITORY_PATH, + ONEDOCKER_CHECKSUM_TYPE, ONEDOCKER_EXE_PATH, ONEDOCKER_REPOSITORY_PATH, ) @@ -59,12 +65,20 @@ "https://one-docker-repository-prod.s3.us-west-2.amazonaws.com/" ) +# The default OneDocker checksum repository path on S3 +DEFAULT_CHECKSUM_REPOSITORY_PATH = ( + "https://one-docker-checksum-repository-stage.s3.us-west-2.amazonaws.com/" +) + # The default path in the Docker image that is going to host the executables DEFAULT_EXE_FOLDER = "/root/onedocker/package/" # the default version of the binary DEFAULT_BINARY_VERSION = "latest" +# the default checksum type for verification +DEFAULT_CHECKSUM_TYPE: str = "SHA256" + logger: logging.Logger @@ -97,16 +111,16 @@ def _prepare_executable( # verify exectuable integrity if checksum_repository_path: _attest_executable( - executable, - checksum_repository_path, - checksum_type, - package_name, - version, + binary_path=executable, + checksum_repository_path=checksum_repository_path, + checksum_type=checksum_type, + package_name=package_name, + version=version, ) else: logger.info("No Checksum Path specified skipping verification") else: - logger.info("Local repository, skip download ...") + logger.info("Local repository, skip download and attestation ...") # grant execute permission to the downloaded executable file if not os.access(executable, os.X_OK): @@ -253,17 +267,28 @@ def _attest_executable( ) logger.info(f"Downloading checksum info for package {package_name}: {version}") - formated_checksum_info = onedocker_checksum_repository.read(package_name, version) - - logger.info(f"Attesting checksum info for package {package_name}: {version}") - attestation_service.attest_binary( - binary_path=binary_path, - package_name=package_name, - version=version, - formated_checksum_info=formated_checksum_info, - checksum_algorithm=checksum_type, - ) - logger.info(f"Finished attestation for package {package_name}: {version}") + try: + formated_checksum_info = onedocker_checksum_repository.read( + package_name, version + ) + + logger.info(f"Attesting checksum info for package {package_name}: {version}") + if formated_checksum_info: + attestation_service.attest_binary( + binary_path=binary_path, + package_name=package_name, + version=version, + formated_checksum_info=formated_checksum_info, + checksum_algorithm=checksum_type, + ) + else: + logger.info("WARNING: No formated ChecksumInfo. Skipping Attestation") + + logger.info(f"Finished attestation for package {package_name}: {version}") + except PcpError: + logger.info( + "WARNING: Connection to StorageService failed. Skipping Attestation" + ) def _parse_package_name(package_name: str) -> str: @@ -312,8 +337,10 @@ def main() -> None: "": str, "--version": str, "--repository_path": schema.Or(None, schema.And(str, len)), + "--checksum_repository_path": schema.Or(None, schema.And(str, len)), "--exe_path": schema.Or(None, schema.And(str, len)), "--exe_args": schema.Or(None, schema.And(str, len)), + "--checksum_type": schema.Or(None, schema.And(str, len)), "--timeout": schema.Or(None, schema.Use(int)), "--log_path": schema.Or(None, schema.Use(Path)), "--cert_params": schema.Or(None, schema.And(str, len)), @@ -335,23 +362,31 @@ def main() -> None: ONEDOCKER_REPOSITORY_PATH, DEFAULT_REPOSITORY_PATH, ) - # [TODO] Update to not be hardcoded - # This will be tied to cli args in a future diff - checksum_repository_path = "" + checksum_repository_path = _read_config( + "checksum_repository_path", + arguments["--checksum_repository_path"], + ONEDOCKER_CHECKSUM_REPOSITORY_PATH, + DEFAULT_CHECKSUM_REPOSITORY_PATH, + ) exe_path = _read_config( "exe_path", arguments["--exe_path"], ONEDOCKER_EXE_PATH, DEFAULT_EXE_FOLDER, ) + checksum_type = ChecksumType( + _read_config( + "checksum_type", + arguments["--checksum_type"], + ONEDOCKER_CHECKSUM_TYPE, + DEFAULT_CHECKSUM_TYPE, + ) + ) certificate_request = ( CertificateRequest.create_instance(arguments["--cert_params"]) if arguments["--cert_params"] else None ) - # [TODO] Update to not be hardoced - # This will be tied to cli args in a future diff - checksum_type = ChecksumType.BLAKE2B _run_package( repository_path=repository_path, diff --git a/onedocker/tests/script/runner/test_onedocker_runner.py b/onedocker/tests/script/runner/test_onedocker_runner.py index a1626479..9bb3ed84 100644 --- a/onedocker/tests/script/runner/test_onedocker_runner.py +++ b/onedocker/tests/script/runner/test_onedocker_runner.py @@ -13,6 +13,7 @@ from fbpcp.entity.certificate_request import CertificateRequest, KeyAlgorithm from fbpcp.error.pcp import InvalidParameterError from onedocker.common.core_dump_handler_aws import AWSCoreDumpHandler +from onedocker.repository.onedocker_checksum import OneDockerChecksumRepository from onedocker.repository.onedocker_package import OneDockerPackageRepository from onedocker.script.runner.onedocker_runner import ( __doc__ as __onedocker_runner_doc__, @@ -136,10 +137,12 @@ def test_main_local_timeout(self): self.assertEqual(cm.exception.code, 1) @patch.object(AttestationService, "attest_binary") + @patch.object(OneDockerChecksumRepository, "read") @patch.object(OneDockerPackageRepository, "download") def test_main( self, mockOneDockerPackageRepositoryDownload, + mockOneDockerChecksumRepositoryRead, mockAttestationServiceVerifyBinary, ): # Arrange @@ -169,12 +172,10 @@ def test_main( "/usr/bin/echo", ) - @patch.object(OneDockerPackageRepository, "download") @patch.object(SelfSignedCertificateService, "generate_certificate") def test_main_good_cert( self, mockSelfSignedCertificateServiceGenerateCertificate, - mockOneDockerPackageRepositoryDownload, ): # Arrange with patch.object( @@ -184,6 +185,7 @@ def test_main_good_cert( "onedocker-runner", "echo", "--version=latest", + "--repository_path=local", "--exe_path=/usr/bin/", "--exe_args=test_message", f"--cert_params={self.test_cert_params}", @@ -196,11 +198,6 @@ def test_main_good_cert( # Assert self.assertEqual(cm.exception.code, 0) mockSelfSignedCertificateServiceGenerateCertificate.assert_called_once_with() - mockOneDockerPackageRepositoryDownload.assert_called_once_with( - "echo", - "latest", - "/usr/bin/echo", - ) def test_main_bad_cert(self): # Arrange From 99d63fd608a398ca345799db95c9a443f5d4812c Mon Sep 17 00:00:00 2001 From: Arnav Sacheti Date: Thu, 14 Jul 2022 11:52:33 -0700 Subject: [PATCH 08/13] Tests for attestation-service (#331) Summary: Pull Request resolved: https://github.com/facebookresearch/fbpcp/pull/331 # Context Now that all the tooling should be done, we want some more comprehensive tests to check everything works as it should # This commit Adds many tests both through s3 and locally to verify that code is working as it should Differential Revision: D36839896 fbshipit-source-id: 70e1905e28b8b538dc99fec1086097072b0badde --- .../script/runner/test_onedocker_runner.py | 104 +++++++++++++++++- 1 file changed, 102 insertions(+), 2 deletions(-) diff --git a/onedocker/tests/script/runner/test_onedocker_runner.py b/onedocker/tests/script/runner/test_onedocker_runner.py index 9bb3ed84..e3fb2623 100644 --- a/onedocker/tests/script/runner/test_onedocker_runner.py +++ b/onedocker/tests/script/runner/test_onedocker_runner.py @@ -13,6 +13,7 @@ from fbpcp.entity.certificate_request import CertificateRequest, KeyAlgorithm from fbpcp.error.pcp import InvalidParameterError from onedocker.common.core_dump_handler_aws import AWSCoreDumpHandler +from onedocker.entity.checksum_type import ChecksumType from onedocker.repository.onedocker_checksum import OneDockerChecksumRepository from onedocker.repository.onedocker_package import OneDockerPackageRepository from onedocker.script.runner.onedocker_runner import ( @@ -143,10 +144,10 @@ def test_main( self, mockOneDockerPackageRepositoryDownload, mockOneDockerChecksumRepositoryRead, - mockAttestationServiceVerifyBinary, + mockAttestationServiceAttestBinary, ): # Arrange - mockAttestationServiceVerifyBinary.return_value = None + mockAttestationServiceAttestBinary.return_value = None with patch.object( sys, "argv", @@ -251,6 +252,105 @@ def test_main_env(self): # Assert self.assertEqual(cm.exception.code, 0) + @patch.object(AttestationService, "attest_binary") + @patch.object(OneDockerChecksumRepository, "read") + @patch.object(OneDockerPackageRepository, "download") + def test_main_checksum( + self, + mockOneDockerPackageRepositoryDownload, + mockOneDockerChecksumRepositoryRead, + mockAttestationServiceAttestBinary, + ): + # Arrange + mockAttestationServiceAttestBinary.return_value = None + mockOneDockerChecksumRepositoryRead.return_value = "checksum_info_goes_here" + with patch.object( + sys, + "argv", + [ + "onedocker-runner", + "ls", + "--version=latest", + "--repository_path=https://onedocker-runner-unittest-asacheti.s3.us-west-2.amazonaws.com/", + "--timeout=1200", + "--exe_path=/usr/bin/", + "--exe_args=-l", + "--checksum_repository_path=https://one-docker-checksum-repository-prod.s3.us-west-2.amazonaws.com/", + "--checksum_type=BLAKE2B", + ], + ): + with self.assertRaises(SystemExit) as cm: + # Act + main() + + # Assert + self.assertEqual(cm.exception.code, 0) + mockOneDockerPackageRepositoryDownload.assert_called_once_with( + "ls", + "latest", + "/usr/bin/ls", + ) + mockAttestationServiceAttestBinary.assert_called_once_with( + binary_path="/usr/bin/ls", + package_name="ls", + version="latest", + formated_checksum_info="checksum_info_goes_here", + checksum_algorithm=ChecksumType.BLAKE2B, + ) + mockOneDockerChecksumRepositoryRead.assert_called_once_with( + "ls", + "latest", + ) + + @patch.object(AttestationService, "attest_binary") + @patch.object(OneDockerChecksumRepository, "read") + @patch.object(OneDockerPackageRepository, "download") + def test_main_checksum_default_type( + self, + mockOneDockerPackageRepositoryDownload, + mockOneDockerChecksumRepositoryRead, + mockAttestationServiceAttestBinary, + ): + # Arrange + mockAttestationServiceAttestBinary.return_value = None + mockOneDockerChecksumRepositoryRead.return_value = "checksum_info_goes_here" + with patch.object( + sys, + "argv", + [ + "onedocker-runner", + "ls", + "--version=latest", + "--repository_path=https://onedocker-runner-unittest-asacheti.s3.us-west-2.amazonaws.com/", + "--timeout=1200", + "--exe_path=/usr/bin/", + "--exe_args=-l", + "--checksum_repository_path=https://one-docker-checksum-repository-prod.s3.us-west-2.amazonaws.com/", + ], + ): + with self.assertRaises(SystemExit) as cm: + # Act + main() + + # Assert + self.assertEqual(cm.exception.code, 0) + mockOneDockerPackageRepositoryDownload.assert_called_once_with( + "ls", + "latest", + "/usr/bin/ls", + ) + mockAttestationServiceAttestBinary.assert_called_once_with( + binary_path="/usr/bin/ls", + package_name="ls", + version="latest", + formated_checksum_info="checksum_info_goes_here", + checksum_algorithm=ChecksumType.SHA256, + ) + mockOneDockerChecksumRepositoryRead.assert_called_once_with( + "ls", + "latest", + ) + @patch("onedocker.script.runner.onedocker_runner.run_cmd") @patch.object(AWSCoreDumpHandler, "locate_core_dump_file") @patch.object(AWSCoreDumpHandler, "upload_core_dump_file") From b95da2f71a0d4206961b57e32b71583036d424e0 Mon Sep 17 00:00:00 2001 From: Arnav Sacheti Date: Thu, 14 Jul 2022 11:52:33 -0700 Subject: [PATCH 09/13] Creating KMS gateway Differential Revision: D37400919 fbshipit-source-id: f52e3b40b40b1d6847c8f414384372a7403c26f0 --- fbpcp/gateway/kms.py | 98 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 98 insertions(+) create mode 100644 fbpcp/gateway/kms.py diff --git a/fbpcp/gateway/kms.py b/fbpcp/gateway/kms.py new file mode 100644 index 00000000..00079446 --- /dev/null +++ b/fbpcp/gateway/kms.py @@ -0,0 +1,98 @@ +#!/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 + +import boto3 +from botocore.client import BaseClient +from fbpcp.decorator.error_handler import error_handler +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 + ) + + @error_handler + def decrypt( + self, + key_id: str, + ciphertext_blob: bytes, + encryption_context: Dict[str, str], + grant_tokens: List[str], + encryption_algorithm: str, + ) -> Dict[str, Any]: + return self.client.encrypt( + KeyId=key_id, + CiphertextBlob=ciphertext_blob, + EncryptionContext=encryption_context, + GrantTokens=grant_tokens, + EncryptionAlgorithm=encryption_algorithm, + ) + + @error_handler + def encrypt( + self, + key_id: str, + plaintext: bytes, + encryption_context: Dict[str, str], + grant_tokens: List[str], + encryption_algorithm: str, + ) -> Dict[str, Any]: + return self.client.encrypt( + KeyId=key_id, + Plaintext=plaintext, + EncryptionContext=encryption_context, + GrantTokens=grant_tokens, + EncryptionAlgorithm=encryption_algorithm, + ) + + @error_handler + def sign( + self, + key_id: str, + message: bytes, + message_type: str, + grant_tokens: List[str], + signing_algorithm: str, + ) -> Dict[str, Any]: + return self.client.sign( + KeyId=key_id, + Message=message, + MessageType=message_type, + GrantTokens=grant_tokens, + SigningAlgorithm=signing_algorithm, + ) + + @error_handler + def verify( + self, + key_id: str, + message: bytes, + message_type: str, + signature: bytes, + signing_algorithm: str, + grant_tokens: List[str], + ) -> Dict[str, Any]: + return self.client.verify( + KeyId=key_id, + Message=message, + MessageType=message_type, + Signature=signature, + SigningAlgorithm=signing_algorithm, + GrantTokens=grant_tokens, + ) From 758d12c6ff8f836c3ff760eedbcc8710286cc5d9 Mon Sep 17 00:00:00 2001 From: Arnav Sacheti Date: Thu, 14 Jul 2022 11:52:33 -0700 Subject: [PATCH 10/13] Creating Key Managment Service for API wrapping Differential Revision: https://www.internalfb.com/diff/D37403863?entry_point=27 fbshipit-source-id: 0647c6f4c042dac2a0d6645dec324609a3b8f668 --- fbpcp/service/key_management.py | 15 ++++++++++ fbpcp/service/key_managment_aws.py | 47 ++++++++++++++++++++++++++++++ 2 files changed, 62 insertions(+) create mode 100644 fbpcp/service/key_management.py create mode 100644 fbpcp/service/key_managment_aws.py diff --git a/fbpcp/service/key_management.py b/fbpcp/service/key_management.py new file mode 100644 index 00000000..843ac62b --- /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 KeyManagmentService(abc.ABC): + @abc.abstractmethod + def sign(self, message: str) -> str: + pass diff --git a/fbpcp/service/key_managment_aws.py b/fbpcp/service/key_managment_aws.py new file mode 100644 index 00000000..df9f59b0 --- /dev/null +++ b/fbpcp/service/key_managment_aws.py @@ -0,0 +1,47 @@ +#!/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 base64 import b64encode +from typing import Any, Dict, List, Optional + +from fbpcp.gateway.kms import KMSGateway +from fbpcp.service.key_management import KeyManagmentService + + +class AWSKeyManagmentService(KeyManagmentService): + 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, + config: Optional[Dict[str, Any]] = None, + ) -> None: + 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") + response = 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, + ) + signature = b64encode(response["Signature"]).decode() + return signature From cf7a080954ed7a9cd127412f628f2d673228acfc Mon Sep 17 00:00:00 2001 From: Arnav Sacheti Date: Thu, 14 Jul 2022 11:52:33 -0700 Subject: [PATCH 11/13] Adding support for plaintext encryption Differential Revision: D37838138 fbshipit-source-id: 7a9e7c044eb1c11b7bafbd451301a956aaba3e85 --- fbpcp/service/key_management.py | 4 ++++ fbpcp/service/key_managment_aws.py | 26 +++++++++++++++++++++++++- 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/fbpcp/service/key_management.py b/fbpcp/service/key_management.py index 843ac62b..a91e95ef 100644 --- a/fbpcp/service/key_management.py +++ b/fbpcp/service/key_management.py @@ -13,3 +13,7 @@ class KeyManagmentService(abc.ABC): @abc.abstractmethod def sign(self, message: str) -> str: pass + + @abc.abstractmethod + def encrypt(self, plaintext: str) -> str: + pass diff --git a/fbpcp/service/key_managment_aws.py b/fbpcp/service/key_managment_aws.py index df9f59b0..3c233bd6 100644 --- a/fbpcp/service/key_managment_aws.py +++ b/fbpcp/service/key_managment_aws.py @@ -15,6 +15,7 @@ class AWSKeyManagmentService(KeyManagmentService): key_id: str + encryption_algorithm: str signing_algorithm: str grant_tokens: List[str] @@ -22,14 +23,22 @@ def __init__( self, region: str, key_id: str, + encryption_algorithm: Optional[str] = None, signing_algorithm: Optional[str] = None, 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: - self.kms_gateway = KMSGateway(region, access_key_id, access_key_data, config) + self.kms_gateway = KMSGateway( + region=region, + access_key_id=access_key_id, + access_key_data=access_key_data, + config=config, + ) self.key_id = key_id + + self.encryption_algorithm = encryption_algorithm if encryption_algorithm else "" self.signing_algorithm = signing_algorithm if signing_algorithm else "" self.grant_tokens = grant_tokens if grant_tokens else [] @@ -45,3 +54,18 @@ def sign(self, message: str, message_type: str = "RAW") -> str: ) signature = b64encode(response["Signature"]).decode() return signature + + def encrypt( + self, plaintext: str, encryption_context: Optional[Dict[str, str]] = None + ) -> str: + if not self.encryption_algorithm: + raise ValueError("No Encryption Algorithm Set") + response = self.kms_gateway.encrypt( + key_id=self.key_id, + plaintext=plaintext, + encryption_context=encryption_context if encryption_context else {}, + grant_tokens=self.grant_tokens, + encryption_algorithm=self.encryption_algorithm, + ) + ciphertext_blob = b64encode(response["CiphertextBlob"]).decode() + return ciphertext_blob From ab97586d18aea9b2e268a97099706d0630e47666 Mon Sep 17 00:00:00 2001 From: Arnav Sacheti Date: Thu, 14 Jul 2022 11:52:33 -0700 Subject: [PATCH 12/13] Adding support for signature verification Differential Revision: D37667144 fbshipit-source-id: 265dd629bd4f9111109b19e8f887d7c733823f21 --- fbpcp/service/key_management.py | 6 +++++- fbpcp/service/key_managment_aws.py | 21 +++++++++++++++++---- 2 files changed, 22 insertions(+), 5 deletions(-) diff --git a/fbpcp/service/key_management.py b/fbpcp/service/key_management.py index a91e95ef..dcb3af56 100644 --- a/fbpcp/service/key_management.py +++ b/fbpcp/service/key_management.py @@ -9,7 +9,7 @@ import abc -class KeyManagmentService(abc.ABC): +class KeyManagementService(abc.ABC): @abc.abstractmethod def sign(self, message: str) -> str: pass @@ -17,3 +17,7 @@ def sign(self, message: str) -> str: @abc.abstractmethod def encrypt(self, plaintext: str) -> str: pass + + @abc.abstractmethod + def verify(self, message: str, signature: str) -> str: + pass diff --git a/fbpcp/service/key_managment_aws.py b/fbpcp/service/key_managment_aws.py index 3c233bd6..dc889129 100644 --- a/fbpcp/service/key_managment_aws.py +++ b/fbpcp/service/key_managment_aws.py @@ -6,18 +6,18 @@ # pyre-strict -from base64 import b64encode +from base64 import b64decode, b64encode from typing import Any, Dict, List, Optional from fbpcp.gateway.kms import KMSGateway -from fbpcp.service.key_management import KeyManagmentService +from fbpcp.service.key_management import KeyManagementService -class AWSKeyManagmentService(KeyManagmentService): +class AWSKeyManagementService(KeyManagementService): key_id: str encryption_algorithm: str signing_algorithm: str - grant_tokens: List[str] + grant_tokens: Optional[List[str]] def __init__( self, @@ -28,6 +28,7 @@ def __init__( grant_tokens: Optional[List[str]] = None, access_key_id: Optional[str] = None, access_key_data: Optional[str] = None, + encryption_context: Optional[Dict[str, str]] = None, config: Optional[Dict[str, Any]] = None, ) -> None: self.kms_gateway = KMSGateway( @@ -69,3 +70,15 @@ def encrypt( ) ciphertext_blob = b64encode(response["CiphertextBlob"]).decode() return ciphertext_blob + + def verify(self, message: str, signature: str, message_type: str = "RAW") -> str: + b64_signature = b64decode(signature.encode()) + response = self.kms_gateway.verify( + key_id=self.key_id, + message=message, + message_type=message_type, + signature=b64_signature, + signing_algorithm=self.signing_algorithm, + grant_tokens=self.grant_tokens, + ) + return response["SignatureValid"] From b8e488216f2caf31bdaac57e6d278f37703d947a Mon Sep 17 00:00:00 2001 From: Arnav Sacheti Date: Thu, 14 Jul 2022 11:57:29 -0700 Subject: [PATCH 13/13] Adding support for message decryption Differential Revision: D37838812 fbshipit-source-id: aebac3c09bff8cbae49fca104b1e659d6b6f173f --- fbpcp/service/key_management.py | 4 ++++ fbpcp/service/key_managment_aws.py | 16 ++++++++++++++++ 2 files changed, 20 insertions(+) diff --git a/fbpcp/service/key_management.py b/fbpcp/service/key_management.py index dcb3af56..54c0a825 100644 --- a/fbpcp/service/key_management.py +++ b/fbpcp/service/key_management.py @@ -14,6 +14,10 @@ class KeyManagementService(abc.ABC): def sign(self, message: str) -> str: pass + @abc.abstractmethod + def decrypt(self, ciphertext_blob: str) -> str: + pass + @abc.abstractmethod def encrypt(self, plaintext: str) -> str: pass diff --git a/fbpcp/service/key_managment_aws.py b/fbpcp/service/key_managment_aws.py index dc889129..b9034658 100644 --- a/fbpcp/service/key_managment_aws.py +++ b/fbpcp/service/key_managment_aws.py @@ -56,6 +56,22 @@ def sign(self, message: str, message_type: str = "RAW") -> str: signature = b64encode(response["Signature"]).decode() return signature + def decrypt( + self, ciphertext_blob: str, encryption_context: Optional[Dict[str, str]] = None + ) -> str: + if not self.encryption_algorithm: + raise ValueError("No Encryption Algorithm Set") + b64_ciphertext_blob = b64decode(ciphertext_blob.encode()) + response = self.kms_gateway.decrypt( + key_id=self.key_id, + ciphertext_blob=b64_ciphertext_blob, + encryption_context=encryption_context, + grant_tokens=self.grant_tokens, + encryption_algorithm=self.encryption_algorithm, + ) + plaintext = b64encode(response["Plaintext"]).decode() + return plaintext + def encrypt( self, plaintext: str, encryption_context: Optional[Dict[str, str]] = None ) -> str: