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, + ) diff --git a/fbpcp/service/key_management.py b/fbpcp/service/key_management.py new file mode 100644 index 00000000..54c0a825 --- /dev/null +++ b/fbpcp/service/key_management.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 + +import abc + + +class KeyManagementService(abc.ABC): + @abc.abstractmethod + 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 + + @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 new file mode 100644 index 00000000..b9034658 --- /dev/null +++ b/fbpcp/service/key_managment_aws.py @@ -0,0 +1,100 @@ +#!/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 b64decode, b64encode +from typing import Any, Dict, List, Optional + +from fbpcp.gateway.kms import KMSGateway +from fbpcp.service.key_management import KeyManagementService + + +class AWSKeyManagementService(KeyManagementService): + key_id: str + encryption_algorithm: str + signing_algorithm: str + grant_tokens: Optional[List[str]] + + 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, + encryption_context: Optional[Dict[str, str]] = None, + config: Optional[Dict[str, Any]] = None, + ) -> None: + 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 [] + + 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 + + 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: + 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 + + 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"] 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/entity/checksum_info.py b/onedocker/entity/checksum_info.py index 66ff933b..22ee575a 100644 --- a/onedocker/entity/checksum_info.py +++ b/onedocker/entity/checksum_info.py @@ -8,10 +8,14 @@ from __future__ import annotations +import base64 +import json + from dataclasses import dataclass -from typing import Any, Dict +from typing import Any, Dict, Optional, Set from onedocker.entity.checksum_type import ChecksumType +from OpenSSL.crypto import FILETYPE_PEM, load_publickey, verify, X509 @dataclass @@ -20,14 +24,16 @@ 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 - Checksums: Dict that holds a pairing between a ChecksumType (Key) and the corresponding hash (Value) + 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) + Signature: Base64 encoded string containing """ package_name: str version: str checksums: Dict[str, str] + signature: str = "" def __post_init__(self) -> None: """ @@ -41,7 +47,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 @@ -63,5 +69,18 @@ def __eq__(self, other: ChecksumInfo) -> bool: return False return True - def asdict(self) -> Dict[str, Any]: - return self.__dict__ + def asdict(self, exclude: Optional[Set[str]] = None) -> Dict[str, Any]: + """ + Returns a dict representation of all fields in ChecksumInfo + + Args: + exclude: Set of Strings that contains all fields to exclude from return + """ + return ( + self.__dict__ + if not exclude + else { + key: self.__dict__[key] for key in self.__dict__ if key not in exclude + } + ) + diff --git a/onedocker/repository/onedocker_checksum.py b/onedocker/repository/onedocker_checksum.py new file mode 100644 index 00000000..04c9612d --- /dev/null +++ b/onedocker/repository/onedocker_checksum.py @@ -0,0 +1,48 @@ +#!/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 _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/script/cli/onedocker_cli.py b/onedocker/script/cli/onedocker_cli.py index 6063e335..96904d33 100644 --- a/onedocker/script/cli/onedocker_cli.py +++ b/onedocker/script/cli/onedocker_cli.py @@ -32,10 +32,12 @@ from docopt import docopt from fbpcp.entity.container_instance import ContainerInstanceStatus from fbpcp.service.container import ContainerService +from fbpcp.service.key_management import KeyManagementService from fbpcp.service.log import LogService 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,8 +45,10 @@ onedocker_svc = None container_svc = None onedocker_package_repo = None -attestation_service = None +onedocker_checksum_repo = None +attestation_svc = None log_svc = None +key_management_svc = None task_definition = None repository_path = None @@ -62,10 +66,29 @@ 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"Signing checksums for package {package_name}: {version}") + checksum_info_signature = key_management_svc.sign( + message=formated_checksum_info, + message_type="RAW", + ) + signed_checksum_info = attestation_svc.add_signature( + formated_checksum_info=formated_checksum_info, + signature=checksum_info_signature, ) - attestation_service.track_binary(package_dir, package_name, version) + + logger.info(f"Uploading checksums for package {package_name}: {version}") + onedocker_checksum_repo.write( + package_name=package_name, + version=version, + checksum_data=signed_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") @@ -160,12 +183,13 @@ 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 _build_key_managment_service(config: Dict[str, Any]) -> KeyManagementService: + key_mangment_class = reflect.get_class(config["class"]) + return key_mangment_class(**config["constructor"]) 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, key_management_svc s = schema.Schema( { "upload": bool, @@ -197,24 +221,28 @@ 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"]) + key_management_svc = _build_key_managment_service( + config["dependency"]["KeyManagementService"] + ) - 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) + 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/config/cli_config_template.yml b/onedocker/script/config/cli_config_template.yml index 560342d2..7d6d3ee1 100644 --- a/onedocker/script/config/cli_config_template.yml +++ b/onedocker/script/config/cli_config_template.yml @@ -12,6 +12,10 @@ onedocker-cli: class: classpath.classname #TODO: change this to actual class name that derived from abstract class: fbpcp.service.log.LogService constructor: attribute_name: value #TODO: change this to actual construction attribute name and value + KeyManagementService: + class: classpath.classname #TODO: change this to actual class name that derived from abstract class: fbpcp.service.key_managment.KeyManagementService + constructor: + attribute_name: value #TODO: change this to actual construction attribute name and value setting: repository_path: repositorypath/ #TODO: change this to the path to your repository, example: https://PATH.s3.REGION.amazonaws.com/ checksum_repository_path: checksumrepositorypath/ #TODO: [OPTIONAL] change this to the path to your repository, example: https://PATH.s3.REGION.amazonaws.com/ diff --git a/onedocker/script/runner/onedocker_runner.py b/onedocker/script/runner/onedocker_runner.py index 25bc6f96..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,16 +40,21 @@ 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, ) 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 @@ -58,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 @@ -96,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): @@ -243,18 +258,37 @@ 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(storage_svc, checksum_repository_path) - - attestation_service.attest_binary( - binary_path=binary_path, - package_name=package_name, - version=version, - checksum_algorithm=checksum_type, + attestation_service = AttestationService() + onedocker_checksum_repository = OneDockerChecksumRepository( + storage_svc, checksum_repository_path ) - logger.info(f"Finished verification for package {package_name}: {version}") + + logger.info(f"Downloading checksum info 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: @@ -303,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)), @@ -326,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/service/attestation.py b/onedocker/service/attestation.py index 82bb160d..82bf901d 100644 --- a/onedocker/service/attestation.py +++ b/onedocker/service/attestation.py @@ -8,9 +8,10 @@ import json import logging -from typing import Any, Dict, List +from typing import Dict, List + +from fbpcp.service.key_management import KeyManagementService -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 @@ -24,19 +25,11 @@ class AttestationService: ChecksumType.SHA256, ChecksumType.BLAKE2B, ] + key_management_svc: KeyManagementService - def __init__(self, storage_svc: StorageService, repository_path: str) -> None: + def __init__(self) -> None: self.logger: logging.Logger = logging.getLogger(__name__) self.checksum_generator = LocalChecksumGenerator() - self.storage_svc = storage_svc - self.repository_path = 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, @@ -56,33 +49,22 @@ 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 Args: - 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 + 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}") @@ -91,20 +73,29 @@ def track_binary( version=version, binary_path=binary_path, ) + formated_checksum_info = json.dumps( + checksum_info.asdict( + exclude={"signature"}, + ), + ) + return formated_checksum_info - # 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(), + def add_signature(self, formated_checksum_info: str, signature: str) -> str: + checksum_dict = json.loads(formated_checksum_info) + checksum_info = ChecksumInfo(**checksum_dict) + checksum_info.signature = signature + + signed_checksum_info = json.dumps( + checksum_info.asdict(), ) + return signed_checksum_info def attest_binary( self, binary_path: str, package_name: str, version: str, + formated_checksum_info: str, checksum_algorithm: ChecksumType, ) -> None: """ @@ -113,32 +104,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/repository/test_onedocker_checksum.py b/onedocker/tests/repository/test_onedocker_checksum.py new file mode 100644 index 00000000..e1eed2b5 --- /dev/null +++ b/onedocker/tests/repository/test_onedocker_checksum.py @@ -0,0 +1,94 @@ +#!/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, + ) + + 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, + ) diff --git a/onedocker/tests/script/cli/test_onedocker_cli.py b/onedocker/tests/script/cli/test_onedocker_cli.py index 00b04610..e0ec0548 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") @@ -365,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 @@ -384,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( diff --git a/onedocker/tests/script/runner/test_onedocker_runner.py b/onedocker/tests/script/runner/test_onedocker_runner.py index 89e4bac3..e3fb2623 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,9 @@ 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.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 ( __doc__ as __onedocker_runner_doc__, @@ -134,14 +138,16 @@ 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, - mockAttestationServiceVerifyBinary, + mockOneDockerChecksumRepositoryRead, + mockAttestationServiceAttestBinary, ): # Arrange - mockAttestationServiceVerifyBinary.return_value = None + mockAttestationServiceAttestBinary.return_value = None with patch.object( sys, "argv", @@ -167,12 +173,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( @@ -182,6 +186,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}", @@ -194,11 +199,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 @@ -252,13 +252,149 @@ 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") + 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 diff --git a/onedocker/tests/service/test_attestation.py b/onedocker/tests/service/test_attestation.py index f5f03b8a..0b64802e 100644 --- a/onedocker/tests/service/test_attestation.py +++ b/onedocker/tests/service/test_attestation.py @@ -6,9 +6,8 @@ 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 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,20 +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) - @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"], + package_name=self.test_package["package_name"], version=self.test_package["version"], ) @@ -72,10 +55,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, @@ -91,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, ) @@ -101,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, @@ -115,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 = { @@ -130,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 @@ -162,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, ) @@ -172,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 @@ -197,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, ) @@ -207,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, @@ -221,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 = { @@ -235,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, ) @@ -245,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"], - )