Add EventSigner implementation

Author: nateprewitt

codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java

@@ -536,6 +536,8 @@ await sleep(retry_token.retry_delay)
 
         writer.pushState(new SignRequestSection());
         if (context.applicationProtocol().isHttpProtocol() && supportsAuth) {
+            writer.addStdlibImport("binascii", "hexlify");
+            writer.addStdlibImport("re");
             writer.write("""
                             # Step 7i: sign the request
                             if auth_option and signer:
@@ -553,6 +555,14 @@ await sleep(retry_token.retry_delay)
                                     )
                                 )
                                 logger.debug("Signed HTTP request: %s", context.transport_request)
+
+                                # TODO - Move this to separate resolution/population function
+                                fields = context._transport_request.fields
+                                auth_value = fields["Authorization"].as_string()  # type: ignore
+                                signature = re.split("Signature=", auth_value)[-1]  # type: ignore
+                                context._properties["signature"] = hexlify(signature.encode('utf-8'))  # type: ignore
+                                context._properties["identity"] = identity
+                                context._properties["signer_properties"] = auth_option.signer_properties
                     """);
         }
         writer.popState();

codegen/core/src/main/java/software/amazon/smithy/python/codegen/integrations/RestJsonProtocolGenerator.java

@@ -426,6 +426,7 @@ public void wrapOutputStream(GenerationContext context, PythonWriter writer) {
                                 transport_response.body  # type: ignore
                             ),
                             deserializer=event_deserializer,  # type: ignore
+                            signer=signer,  # type: ignore
                         )
                         """);
     }

packages/aws-event-stream/src/aws_event_stream/_private/serializers.py

@@ -2,10 +2,10 @@
 # SPDX-License-Identifier: Apache-2.0
 import asyncio
 import datetime
-from collections.abc import Callable, Iterator
+from collections.abc import Iterator
 from contextlib import contextmanager
 from io import BytesIO
-from typing import Never
+from typing import Never, Protocol
 
 from smithy_core.aio.interfaces import AsyncWriter
 from smithy_core.codecs import Codec
@@ -20,7 +20,7 @@
 from smithy_core.shapes import ShapeType
 from smithy_event_stream.aio.interfaces import AsyncEventPublisher
 
-from ..events import EventMessage, HEADER_VALUE, Short, Byte, Long
+from ..events import EventHeaderEncoder, EventMessage, HEADER_VALUE, Short, Byte, Long
 from ..exceptions import InvalidHeaderValue
 from . import (
     INITIAL_REQUEST_EVENT_TYPE,
@@ -33,16 +33,23 @@
 _DEFAULT_BLOB_CONTENT_TYPE = "application/octet-stream"
 
 
-type Signer = Callable[[EventMessage], EventMessage]
-"""A function that takes an event message and signs it, and returns it signed."""
+class EventSigner(Protocol):
+    """A signer to manage credentials and EventMessages for an Event Stream lifecyle."""
+
+    def sign_event(
+        self,
+        *,
+        event_message: EventMessage,
+        event_encoder_cls: type[EventHeaderEncoder],
+    ) -> EventMessage: ...
 
 
 class AWSAsyncEventPublisher[E: SerializeableShape](AsyncEventPublisher[E]):
     def __init__(
         self,
         payload_codec: Codec,
         async_writer: AsyncWriter,
-        signer: Signer | None = None,
+        signer: EventSigner | None = None,
         is_client_mode: bool = True,
     ):
         self._writer = async_writer
@@ -62,8 +69,11 @@ async def send(self, event: E) -> None:
                 "Expected an event message to be serialized, but was None."
             )
         if self._signer is not None:
-            result = self._signer(result)
-
+            encoder = self._serializer.event_header_encoder_cls
+            result = await self._signer.sign_event(
+                event_message=result,
+                event_encoder_cls=encoder,
+            )
         encoded_result = result.encode()
         try:
             await self._writer.write(encoded_result)
@@ -97,6 +107,7 @@ def __init__(
             self._initial_message_event_type = INITIAL_REQUEST_EVENT_TYPE
         else:
             self._initial_message_event_type = INITIAL_RESPONSE_EVENT_TYPE
+        self.event_header_encoder_cls = EventHeaderEncoder
 
     def get_result(self) -> EventMessage | None:
         return self._result

packages/aws-event-stream/src/aws_event_stream/aio/__init__.py

@@ -17,7 +17,7 @@
 
 from .._private.deserializers import AWSAsyncEventReceiver as _AWSEventReceiver
 from .._private.serializers import AWSAsyncEventPublisher as _AWSEventPublisher
-from .._private.serializers import Signer
+from .._private.serializers import EventSigner
 from ..exceptions import MissingInitialResponse
 
 
@@ -36,7 +36,7 @@ def __init__(
         awaitable_response: Awaitable[Response],
         awaitable_output: Awaitable[R],
         deserializeable_response: type[R] | None = None,
-        signer: Signer | None = None,
+        signer: EventSigner | None = None,
         is_client_mode: bool = True,
     ) -> None:
         """Construct an AWSDuplexEventStream.
@@ -116,7 +116,7 @@ def __init__(
         payload_codec: Codec,
         async_writer: AsyncWriter,
         awaitable_output: Awaitable[R],
-        signer: Signer | None = None,
+        signer: EventSigner | None = None,
         is_client_mode: bool = True,
     ) -> None:
         """Construct an AWSInputEventStream.

packages/aws-sdk-signers/src/aws_sdk_signers/__init__.py

@@ -9,14 +9,20 @@
 from ._http import URI, AWSRequest, Field, Fields
 from ._identity import AWSCredentialIdentity
 from ._io import AsyncBytesReader
-from .signers import AsyncSigV4Signer, SigV4Signer, SigV4SigningProperties
+from .signers import (
+    AsyncSigV4Signer,
+    AsyncEventSigner,
+    SigV4Signer,
+    SigV4SigningProperties,
+)
 
 __license__ = "Apache-2.0"
 __version__ = importlib.metadata.version("aws-sdk-signers")
 
 __all__ = (
     "AsyncBytesReader",
     "AsyncSigV4Signer",
+    "AsyncEventSigner",
     "AWSCredentialIdentity",
     "AWSRequest",
     "Field",

packages/aws-sdk-signers/src/aws_sdk_signers/interfaces/events.py

@@ -0,0 +1,47 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import datetime
+import uuid
+from collections.abc import Mapping
+from typing import Protocol
+
+
+type HEADER_VALUE = bool | int | bytes | str | datetime.datetime | uuid.UUID
+"""A union of valid value types for event headers."""
+
+
+type HEADERS_DICT = Mapping[str, HEADER_VALUE]
+"""A dictionary of event headers."""
+
+
+class EventMessage(Protocol):
+    """A signable message that may be sent over an event stream."""
+
+    headers: HEADERS_DICT
+    """The headers present in the event message."""
+
+    payload: bytes
+    """The serialized bytes of the message payload."""
+
+    def encode(self) -> bytes:
+        """Encode heads and payload into bytes for transit."""
+        ...
+
+
+class EventHeaderEncoder(Protocol):
+    """A utility class that encodes event headers into bytes."""
+
+    def clear(self) -> None:
+        """Clear all previously encoded headers."""
+        ...
+
+    def get_result(self) -> bytes:
+        """Get all the encoded header bytes."""
+        ...
+
+    def encode_headers(self, headers: HEADERS_DICT) -> None:
+        """Encode a map of headers."""
+        ...

packages/aws-sdk-signers/src/aws_sdk_signers/signers.py

@@ -1,15 +1,17 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+import asyncio
 import datetime
 import hmac
 import io
 import warnings
 from asyncio import iscoroutinefunction
+from binascii import hexlify
 from collections.abc import AsyncIterable, Iterable
 from copy import deepcopy
 from hashlib import sha256
-from typing import Required, TypedDict
+from typing import Required, TypedDict, TYPE_CHECKING
 from urllib.parse import parse_qsl, quote
 
 from .interfaces.io import AsyncSeekable, Seekable
@@ -19,6 +21,9 @@
 from ._io import AsyncBytesReader
 from .exceptions import AWSSDKWarning, MissingExpectedParameterException
 
+if TYPE_CHECKING:
+    from .interfaces.events import EventMessage, EventHeaderEncoder
+
 HEADERS_EXCLUDED_FROM_SIGNING: tuple[str, ...] = (
     "accept",
     "accept-encoding",
@@ -789,6 +794,108 @@ async def _compute_payload_hash(
         return checksum.hexdigest()
 
 
+class AsyncEventSigner:
+    def __init__(
+        self,
+        *,
+        signing_properties: SigV4SigningProperties,
+        identity: AWSCredentialIdentity,
+        initial_signature: bytes,
+    ):
+        self._signing_properties = signing_properties
+        self._identity = identity
+        self._prior_signature = initial_signature
+        self._signing_lock = asyncio.Lock()
+
+    async def sign_event(
+        self,
+        *,
+        event_message: "EventMessage",
+        event_encoder_cls: type["EventHeaderEncoder"],
+    ) -> "EventMessage":
+        async with self._signing_lock:
+            # Copy and prepopulate any missing values in the
+            # signing properties.
+            new_signing_properties = SigV4SigningProperties(  # type: ignore
+                **self._signing_properties
+            )
+            if "date" not in new_signing_properties:
+                date_obj = datetime.datetime.now(datetime.UTC)
+                new_signing_properties["date"] = date_obj.strftime(
+                    SIGV4_TIMESTAMP_FORMAT
+                )
+
+            timestamp = new_signing_properties["date"]
+            headers: dict[str, str | bytes] = {":date": timestamp}
+            encoder = event_encoder_cls()
+            encoder.encode_headers(event_message.headers)
+            encoded_headers = encoder.get_result()
+
+            string_to_sign = await self._event_string_to_sign(
+                timestamp=timestamp,
+                scope=self._scope(new_signing_properties),
+                encoded_headers=encoded_headers,
+                payload=event_message.payload,
+                prior_signature=self._prior_signature,
+            )
+            event_signature = await self._sign_event(
+                timestamp=timestamp,
+                string_to_sign=string_to_sign,
+                signing_properties=new_signing_properties,
+            )
+            headers[":chunk-signature"] = event_signature
+            event_message.headers.update(headers)  # type: ignore
+
+            # set new prior signature before releasing the lock
+            self._prior_signature = event_signature
+
+        return event_message
+
+    async def _event_string_to_sign(
+        self,
+        *,
+        timestamp: str,
+        scope: str,
+        encoded_headers: bytes,
+        payload: bytes,
+        prior_signature: bytes,
+    ) -> str:
+        return (
+            "AWS-HMAC-SHA256-PAYLOAD\n"
+            f"{timestamp}\n"
+            f"{scope}\n"
+            f"{hexlify(prior_signature).decode('utf-8')}\n"
+            f"{sha256(encoded_headers).hexdigest()}\n"
+            f"{sha256(payload).hexdigest()}"
+        )
+
+    async def _sign_event(
+        self,
+        *,
+        timestamp: str,
+        string_to_sign: str,
+        signing_properties: SigV4SigningProperties,
+    ) -> bytes:
+        key = self._identity.secret_access_key.encode("utf-8")
+        today = timestamp[:8].encode("utf-8")
+        k_date = self._hash(b"AWS4" + key, today)
+        k_region = self._hash(k_date, signing_properties["region"].encode("utf-8"))
+        k_service = self._hash(k_region, signing_properties["service"].encode("utf-8"))
+        k_signing = self._hash(k_service, b"aws4_request")
+        return self._hash(k_signing, string_to_sign.encode("utf-8"))
+
+    def _hash(self, key: bytes, msg: bytes) -> bytes:
+        return hmac.new(key, msg, sha256).digest()
+
+    def _scope(self, signing_properties: SigV4SigningProperties) -> str:
+        assert "date" in signing_properties
+        formatted_date = signing_properties["date"][0:8]
+        region = signing_properties["region"]
+        service = signing_properties["service"]
+        # Scope format: <YYYYMMDD>/<AWS Region>/<AWS Service>/aws4_request
+        return f"{formatted_date}/{region}/{service}/aws4_request"
+
+
 def _remove_dot_segments(path: str, remove_consecutive_slashes: bool = True) -> str:
     """Removes dot segments from a path per :rfc:`3986#section-5.2.4`.