feat: add PEP 561 py.typed marker to support mypy type analysis (#563)

MANIFEST.in

@@ -3,6 +3,7 @@ include CHANGELOG.rst
 include CONTRIBUTING.rst
 include LICENSE
 include requirements.txt
+recursive-include src/aws_encryption_sdk *.pyi py.typed
 
 recursive-include doc *
 recursive-include test *.py

setup.py

@@ -32,6 +32,7 @@ def get_requirements():
     name="aws-encryption-sdk",
     packages=find_packages("src"),
     package_dir={"": "src"},
+    package_data={"aws_encryption_sdk": ["py.typed", "*.pyi", "*/*.pyi"]},
     version=get_version(),
     author="Amazon Web Services",
     maintainer="Amazon Web Services",

src/aws_encryption_sdk/__init__.pyi

@@ -0,0 +1,132 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+import sys
+from typing import Any, BinaryIO, Dict, IO, Optional, Tuple, Union, overload
+
+if sys.version_info >= (3, 8):
+    from typing import Literal
+else:
+    from typing_extensions import Literal
+
+from aws_encryption_sdk.caches.local import LocalCryptoMaterialsCache
+from aws_encryption_sdk.caches.null import NullCryptoMaterialsCache
+from aws_encryption_sdk.exceptions import AWSEncryptionSDKClientError
+from aws_encryption_sdk.identifiers import Algorithm, CommitmentPolicy, __version__
+from aws_encryption_sdk.internal.utils.signature import SignaturePolicy
+from aws_encryption_sdk.key_providers.kms import (
+    DiscoveryAwsKmsMasterKeyProvider,
+    KMSMasterKeyProviderConfig,
+    StrictAwsKmsMasterKeyProvider,
+)
+from aws_encryption_sdk.materials_managers.base import CryptoMaterialsManager
+from aws_encryption_sdk.materials_managers.caching import CachingCryptoMaterialsManager
+from aws_encryption_sdk.materials_managers.default import DefaultCryptoMaterialsManager
+from aws_encryption_sdk.streaming_client import (
+    DecryptorConfig,
+    EncryptorConfig,
+    StreamDecryptor,
+    StreamEncryptor,
+)
+from aws_encryption_sdk.structures import MessageHeader
+
+__all__ = [
+    "LocalCryptoMaterialsCache",
+    "NullCryptoMaterialsCache",
+    "AWSEncryptionSDKClientError",
+    "Algorithm",
+    "CommitmentPolicy",
+    "SignaturePolicy",
+    "DiscoveryAwsKmsMasterKeyProvider",
+    "KMSMasterKeyProviderConfig",
+    "StrictAwsKmsMasterKeyProvider",
+    "CachingCryptoMaterialsManager",
+    "DefaultCryptoMaterialsManager",
+    "DecryptorConfig",
+    "EncryptorConfig",
+    "StreamDecryptor",
+    "StreamEncryptor",
+    "MessageHeader",
+    "EncryptionSDKClientConfig",
+    "EncryptionSDKClient",
+    "__version__",
+]
+
+class EncryptionSDKClientConfig:
+    commitment_policy: CommitmentPolicy
+    max_encrypted_data_keys: Optional[int]
+    def __init__(
+        self,
+        commitment_policy: CommitmentPolicy = ...,
+        max_encrypted_data_keys: Optional[int] = ...,
+    ) -> None: ...
+
+class EncryptionSDKClient:
+    config: EncryptionSDKClientConfig
+    def __init__(
+        self,
+        config: Optional[EncryptionSDKClientConfig] = ...,
+        commitment_policy: Optional[CommitmentPolicy] = ...,
+        max_encrypted_data_keys: Optional[int] = ...,
+        **kwargs: Any,
+    ) -> None: ...
+    def encrypt(
+        self,
+        source: Union[str, bytes, IO[bytes], BinaryIO],
+        materials_manager: Optional[CryptoMaterialsManager] = ...,
+        key_provider: Optional[Any] = ...,
+        keyring: Optional[Any] = ...,
+        source_length: Optional[int] = ...,
+        encryption_context: Optional[Dict[str, str]] = ...,
+        algorithm: Optional[Algorithm] = ...,
+        frame_length: Optional[int] = ...,
+        config: Optional[EncryptorConfig] = ...,
+        **kwargs: Any,
+    ) -> Tuple[bytes, MessageHeader]: ...
+    def decrypt(
+        self,
+        source: Union[str, bytes, IO[bytes], BinaryIO],
+        materials_manager: Optional[CryptoMaterialsManager] = ...,
+        key_provider: Optional[Any] = ...,
+        keyring: Optional[Any] = ...,
+        source_length: Optional[int] = ...,
+        encryption_context: Optional[Dict[str, str]] = ...,
+        max_body_length: Optional[int] = ...,
+        config: Optional[DecryptorConfig] = ...,
+        **kwargs: Any,
+    ) -> Tuple[bytes, MessageHeader]: ...
+    @overload
+    def stream(
+        self,
+        mode: Literal["e", "encrypt"],
+        source: Optional[Union[str, bytes, IO[bytes], BinaryIO]] = ...,
+        materials_manager: Optional[CryptoMaterialsManager] = ...,
+        key_provider: Optional[Any] = ...,
+        keyring: Optional[Any] = ...,
+        source_length: Optional[int] = ...,
+        encryption_context: Optional[Dict[str, str]] = ...,
+        algorithm: Optional[Algorithm] = ...,
+        frame_length: Optional[int] = ...,
+        config: Optional[EncryptorConfig] = ...,
+        **kwargs: Any,
+    ) -> StreamEncryptor: ...
+    @overload
+    def stream(
+        self,
+        mode: Literal["d", "decrypt", "decrypt-unsigned"],
+        source: Optional[Union[str, bytes, IO[bytes], BinaryIO]] = ...,
+        materials_manager: Optional[CryptoMaterialsManager] = ...,
+        key_provider: Optional[Any] = ...,
+        keyring: Optional[Any] = ...,
+        source_length: Optional[int] = ...,
+        encryption_context: Optional[Dict[str, str]] = ...,
+        max_body_length: Optional[int] = ...,
+        config: Optional[DecryptorConfig] = ...,
+        **kwargs: Any,
+    ) -> StreamDecryptor: ...
+    @overload
+    def stream(
+        self,
+        mode: str,
+        source: Optional[Union[str, bytes, IO[bytes], BinaryIO]] = ...,
+        **kwargs: Any,
+    ) -> Union[StreamEncryptor, StreamDecryptor]: ...

src/aws_encryption_sdk/exceptions.pyi

@@ -0,0 +1,18 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+class AWSEncryptionSDKClientError(Exception): ...
+class BadCiphertextError(AWSEncryptionSDKClientError): ...
+class DecryptKeyError(AWSEncryptionSDKClientError): ...
+class EncryptKeyError(AWSEncryptionSDKClientError): ...
+class MasterKeyProviderError(AWSEncryptionSDKClientError): ...
+class NotFoundError(AWSEncryptionSDKClientError): ...
+class NotSupportedError(AWSEncryptionSDKClientError): ...
+class SerializationError(AWSEncryptionSDKClientError): ...
+class UnknownIdentityError(AWSEncryptionSDKClientError): ...
+class ActionNotPermittedError(AWSEncryptionSDKClientError): ...
+class InvalidAlgorithmError(AWSEncryptionSDKClientError): ...
+class InvalidConfigError(AWSEncryptionSDKClientError): ...
+class GenerateKeyError(AWSEncryptionSDKClientError): ...
+class CustomKeyStoreError(AWSEncryptionSDKClientError): ...
+class CacheError(AWSEncryptionSDKClientError): ...
+class MaxConsecutiveSleepRoundsExceededError(AWSEncryptionSDKClientError): ...

src/aws_encryption_sdk/identifiers.pyi

@@ -0,0 +1,108 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+from enum import Enum
+from typing import Any, Dict, Optional, Tuple, Type, TypeVar
+
+__version__: str
+USER_AGENT_SUFFIX: str
+
+class SerializationVersion(Enum):
+    V1 = ...
+    V2 = ...
+
+class ObjectType(Enum):
+    CUSTOMER_AE_DATA = ...
+
+class ContentType(Enum):
+    NO_FRAMING = ...
+    FRAMED_DATA = ...
+
+class ContentAADString(Enum):
+    FRAME_STRING_ID = ...
+    FINAL_FRAME_STRING_ID = ...
+    NON_FRAMED_STRING_ID = ...
+
+class CommitmentPolicy(Enum):
+    FORBID_ENCRYPT_ALLOW_DECRYPT = ...
+    REQUIRE_ENCRYPT_ALLOW_DECRYPT = ...
+    REQUIRE_ENCRYPT_REQUIRE_DECRYPT = ...
+
+class EncryptionType(Enum):
+    SYMMETRIC = ...
+    ASYMMETRIC = ...
+
+class EncryptionKeyType(Enum):
+    SYMMETRIC = ...
+    RSA = ...
+
+class EncryptionSuite(Enum):
+    AES_128_GCM_IV12_TAG16 = ...
+    AES_192_GCM_IV12_TAG16 = ...
+    AES_256_GCM_IV12_TAG16 = ...
+    algorithm: Any
+    mode: Any
+    data_key_length: int
+    iv_length: int
+    auth_length: int
+    tag_len: int
+    auth_key_length: int
+
+class KDFSuite(Enum): ...
+class AuthenticationSuite(Enum): ...
+
+_T = TypeVar("_T", bound="AlgorithmSuite")
+
+class AlgorithmSuite(Enum):
+    AES_128_GCM_IV12_TAG16 = ...
+    AES_192_GCM_IV12_TAG16 = ...
+    AES_256_GCM_IV12_TAG16 = ...
+    AES_128_GCM_IV12_TAG16_HKDF_SHA256 = ...
+    AES_192_GCM_IV12_TAG16_HKDF_SHA256 = ...
+    AES_256_GCM_IV12_TAG16_HKDF_SHA256 = ...
+    AES_128_GCM_IV12_TAG16_HKDF_SHA256_ECDSA_P256 = ...
+    AES_192_GCM_IV12_TAG16_HKDF_SHA384_ECDSA_P384 = ...
+    AES_256_GCM_IV12_TAG16_HKDF_SHA384_ECDSA_P384 = ...
+    AES_256_GCM_HKDF_SHA512_COMMIT_KEY = ...
+    AES_256_GCM_HKDF_SHA512_COMMIT_KEY_ECDSA_P384 = ...
+
+    algorithm_id: int
+    encryption: Any
+    message_format_version: SerializationVersion
+    kdf: Any
+    authentication: Any
+    allowed: bool
+    encryption_algorithm: Any
+    encryption_mode: Any
+    data_key_len: int
+    iv_len: int
+    auth_key_len: int
+    auth_len: int
+    kdf_type: Any
+    kdf_hash_type: Any
+    signing_algorithm_info: Any
+    signing_hash_type: Any
+    signature_len: int
+    header_auth_iv: bytes
+
+    @classmethod
+    def get_by_id(cls: Type[_T], algorithm_id: int) -> _T: ...
+    def id_as_bytes(self) -> bytes: ...
+    def is_cacheable(self) -> bool: ...
+    def is_committing(self) -> bool: ...
+    def is_signing(self) -> bool: ...
+    def algorithm_suite_data_length(self) -> int: ...
+
+Algorithm = AlgorithmSuite
+
+class WrappingAlgorithm(Enum):
+    AES_128_GCM_IV12_TAG16_NO_PADDING = ...
+    AES_192_GCM_IV12_TAG16_NO_PADDING = ...
+    AES_256_GCM_IV12_TAG16_NO_PADDING = ...
+    RSA_OAEP_SHA1_MGF1 = ...
+    RSA_OAEP_SHA256_MGF1 = ...
+    RSA_OAEP_SHA384_MGF1 = ...
+    RSA_OAEP_SHA512_MGF1 = ...
+    RSA_PKCS1 = ...
+
+class SequenceIdentifier(Enum):
+    SEQUENCE_NUMBER = ...

src/aws_encryption_sdk/materials_managers/base.pyi

@@ -0,0 +1,10 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+import abc
+from typing import Any
+
+class CryptoMaterialsManager(abc.ABC):
+    @abc.abstractmethod
+    def get_encryption_materials(self, request: Any) -> Any: ...
+    @abc.abstractmethod
+    def decrypt_materials(self, request: Any) -> Any: ...

src/aws_encryption_sdk/materials_managers/caching.pyi

@@ -0,0 +1,17 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+from typing import Any, Optional
+from aws_encryption_sdk.materials_managers.base import CryptoMaterialsManager
+
+class CachingCryptoMaterialsManager(CryptoMaterialsManager):
+    def __init__(
+        self,
+        cache: Any,
+        master_key_provider: Optional[Any] = ...,
+        backing_materials_manager: Optional[CryptoMaterialsManager] = ...,
+        max_age: Optional[float] = ...,
+        max_messages_encrypted: Optional[int] = ...,
+        max_bytes_encrypted: Optional[int] = ...,
+    ) -> None: ...
+    def get_encryption_materials(self, request: Any) -> Any: ...
+    def decrypt_materials(self, request: Any) -> Any: ...

src/aws_encryption_sdk/materials_managers/default.pyi

@@ -0,0 +1,9 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+from typing import Any, Optional
+from aws_encryption_sdk.materials_managers.base import CryptoMaterialsManager
+
+class DefaultCryptoMaterialsManager(CryptoMaterialsManager):
+    def __init__(self, master_key_provider: Optional[Any] = ...) -> None: ...
+    def get_encryption_materials(self, request: Any) -> Any: ...
+    def decrypt_materials(self, request: Any) -> Any: ...

src/aws_encryption_sdk/py.typed

@@ -0,0 +1 @@
+partial

src/aws_encryption_sdk/streaming_client.pyi

@@ -0,0 +1,76 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+import types
+from typing import Any, BinaryIO, Dict, IO, Iterator, Optional, Type, Union
+from aws_encryption_sdk.identifiers import Algorithm, CommitmentPolicy
+from aws_encryption_sdk.materials_managers.base import CryptoMaterialsManager
+from aws_encryption_sdk.structures import MessageHeader
+
+class _ClientConfig:
+    materials_manager: CryptoMaterialsManager
+    source: Any
+    source_length: Optional[int]
+    def __init__(
+        self,
+        materials_manager: Optional[CryptoMaterialsManager] = ...,
+        source: Optional[Union[str, bytes, IO[bytes], BinaryIO]] = ...,
+        source_length: Optional[int] = ...,
+        key_provider: Optional[Any] = ...,
+        keyring: Optional[Any] = ...,
+        **kwargs: Any,
+    ) -> None: ...
+
+class EncryptorConfig(_ClientConfig):
+    encryption_context: Dict[str, str]
+    algorithm: Algorithm
+    frame_length: int
+    commitment_policy: CommitmentPolicy
+    max_encrypted_data_keys: Optional[int]
+    def __init__(
+        self,
+        encryption_context: Optional[Dict[str, str]] = ...,
+        algorithm: Optional[Algorithm] = ...,
+        frame_length: Optional[int] = ...,
+        commitment_policy: Optional[CommitmentPolicy] = ...,
+        max_encrypted_data_keys: Optional[int] = ...,
+        **kwargs: Any,
+    ) -> None: ...
+
+class DecryptorConfig(_ClientConfig):
+    max_body_length: Optional[int]
+    commitment_policy: CommitmentPolicy
+    max_encrypted_data_keys: Optional[int]
+    encryption_context: Optional[Dict[str, str]]
+    def __init__(
+        self,
+        max_body_length: Optional[int] = ...,
+        commitment_policy: Optional[CommitmentPolicy] = ...,
+        max_encrypted_data_keys: Optional[int] = ...,
+        encryption_context: Optional[Dict[str, str]] = ...,
+        **kwargs: Any,
+    ) -> None: ...
+
+class _EncryptionStream:
+    header: MessageHeader
+    config: _ClientConfig
+    def __init__(self, **kwargs: Any) -> None: ...
+    def read(self, size: Optional[int] = ...) -> bytes: ...
+    def close(self) -> None: ...
+    def closed(self) -> bool: ...
+    def __enter__(self) -> _EncryptionStream: ...
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[types.TracebackType],
+    ) -> None: ...
+    def __iter__(self) -> Iterator[bytes]: ...
+    def __next__(self) -> bytes: ...
+
+class StreamEncryptor(_EncryptionStream):
+    config: EncryptorConfig
+    def __enter__(self) -> StreamEncryptor: ...
+
+class StreamDecryptor(_EncryptionStream):
+    config: DecryptorConfig
+    def __enter__(self) -> StreamDecryptor: ...