clear-skies-aws 1.10.2__py3-none-any.whl → 2.0.2__py3-none-any.whl

clearskies_aws/mocks/actions/step_function.py

@@ -1,26 +1,32 @@
-from types import ModuleType
+from __future__ import annotations
+
+from typing import Any
+
 from clearskies import Model
-from ...actions.step_function import StepFunction as BaseStepFunction
-class StepFunction(BaseStepFunction):
-    calls = None
+from types_boto3_stepfunctions import SFNClient
 
-    def __init__(self, environment, boto3, di):
-        super().__init__(environment, boto3, di)
+from clearskies_aws.actions.step_function import StepFunction as BaseStepFunction
+
+
+class StepFunction(BaseStepFunction):
+    calls: list[dict[str, Any]] | None = None
 
     @classmethod
     def mock(cls, di):
-        StepFunction.calls = []
+        cls.calls = []
         di.mock_class(BaseStepFunction, StepFunction)
 
-    def _execute_action(self, client: ModuleType, model: Model) -> None:
+    def _execute_action(self, client: SFNClient, model: Model) -> None:
         """Send a notification as configured."""
-        if StepFunction.calls == None:
+        if StepFunction.calls is None:
             StepFunction.calls = []
 
-        StepFunction.calls.append({
-            "stateMachineArn": self.get_arn(model),
-            "input": self.get_message_body(model),
-        })
+        StepFunction.calls.append(
+            {
+                "stateMachineArn": self.get_arn(model),
+                "input": self.get_message_body(model),
+            }
+        )
 
         if self.column_to_store_execution_arn:
             model.save({self.column_to_store_execution_arn: "mock_execution_arn"})

clearskies_aws/models/web_socket_connection_model.py

@@ -0,0 +1,182 @@
+from __future__ import annotations
+
+import json
+
+import clearskies
+
+import clearskies_aws
+
+
+class WebSocketConnectionModel(clearskies.Model):
+    """
+    Help manage message sending to websocket connections.
+
+    ## Working with Websockets
+
+    This is a partial model class to help send messages to websocket connections in an API gateway.
+    With a API Gateway managed websocket, the API gateway assigns an id to every connection, and you
+    send messages to the API gateway itself flagged for some client, via its connection id.  It
+    helps to understand that you don't need to be connected to the websocket itself to send messages
+    to the things connected to it.  Instead, you just need the necessary permission on the API gateway
+    and you need to know the connection id of the client you want to send a message to.  For reference,
+    the necessary AWS permission is:
+
+    ```
+    {
+        "Version": "2012-10-17",
+        "Statement": [
+            {
+                "Effect": "Allow",
+                "Action": ["execute-api:ManageConnections"],
+                "Resource": "arn:aws:execute-api:${aws_region_name}:${aws_account_id}:${api_gateway_id}/{stage}/*",
+            }
+        ],
+    }
+    ```
+
+    A simple flow that reproduces a pub/sub approach is:
+
+     1. Client connects to the API gateway via a websocket, and the backend records the connection id
+        in a backend somewhere
+     2. Client sends a message through the websocket to "register"/"subscribe" for some resource, and
+         the backend service updates the record for the connection to record what resource it is
+         subscribed to
+     3. If a server needs to send a message to everyone subscribed to a resource, it queries the backends
+        for all records connected to the resource in question and sends a message to their connection ids
+        through the backend.
+     4. If a client needs to send a message to everyone subscribed to a resource, it sends a message
+        through the websocket to some backend service which then passes the message along to the appropriate
+        connections, just as in #3 above.
+
+    Most examples with Websockets and API Gateway use dynamodb for storage.  You can, of course, use whatever
+    backend you want.  Still, below is an example pub/sub application to demonstrate how to build a basic
+    websocket app:
+
+    ```
+    import clearskies
+    import clearskies_aws
+
+    #####################
+    ## Our model class ##
+    #####################
+
+    class Client(clearskies_aws.models.WebSocketConnectionModel):
+        backend = clearskies_aws.backends.DynamoDbBackend()
+
+        # the base WebSocketConnectionModel class defines a string
+        # column called `connection_id` and sets it as the id column,
+        # so those are already set.  We just need any additional columns
+        resource_id = clearskies.columns.String()
+
+    ###########################
+    ## Our application logic ##
+    ###########################
+
+    def on_connect(clients, connection_id):
+        clients.create({"connection_id": connection_id})
+
+    def on_subscribe(clients, connection_id, request_data):
+        client = clients.find(f"connection_id={connection_id}")
+        if client.exists:
+            # we blindly save the request id, which makes it user-generated.  This also
+            # allows the client to "unsubscribe" by sending up a blank resource
+            client.save({"resource_id": request_data["resource_id"])
+
+    def on_publish(clients, connection_id, request_data):
+        my_client = clients.find(f"connection_id={connection_id}")
+        message = request_data.get("message")
+
+        # The problem with our standard input validation is that we can't return a response
+        # to the client with an error message.  This is not a transactional system.  Instead,
+        # we would have to send the client a new message with the error in it, which is not
+        # something the clearskies endpoints are designed for.
+        if not message or not my_client.resource_id:
+            return
+
+        for client in clients.where(f"resource_id={my_client.resource_id}").paginate_all():
+            if client.connection_id == my_client.connection_id:
+                continue
+
+            # the send function is provided by clearskies_aws.models.WebSocketConnectionModel
+            client.send({"message": message})
+
+    def on_disconnect(clients, connection_id):
+        clients.find(f"connection_id={connection_id}").delete(except_if_not_exists=False)
+
+    ######################################
+    ## Wiring it all up with clearskies ##
+    ######################################
+
+    # We're going to build one application, even though each action gets it's own lambda.
+    # The URLs aren't used for routing, but simply to allow us to select which function
+    # is associated with each lambda.  Actual routing still happens in the API Gateway
+    websocket_application = clearskies_aws.contexts.LambdaApiGatewayWebSocket(
+        clearskies.EndpointGroup([
+            clearskies.endpoints.Callable(
+                on_connect,
+                url="on_connect",
+            ),
+            clearskies.endpoints.Callable(
+                on_subscribe,
+                url="on_subscribe",
+            ),
+            clearskies.endpoints.Callable(
+                on_publish,
+                url="on_publish",
+            ),
+            clearskies.endpoints.Callable(
+                on_disconnect,
+                url="on_disconnect",
+            ),
+        ]),
+        classes=[Client],
+    )
+
+    ################################
+    ## The actual lambda handlers ##
+    ################################
+
+    def on_connect_handler(event, context):
+        return websocket_application(url="on_connect")
+
+    def on_subscribe_handler(event, context):
+        return websocket_application(url="on_subscribe")
+
+    def on_publish_handler(event, context):
+        return websocket_application(url="on_publish")
+
+    def on_disconnect_handler(event, context):
+        return websocket_application(url="on_disconnect")
+    ```
+
+    """
+
+    id_column_name = "connection_id"
+
+    boto3 = clearskies_aws.di.inject.Boto3()
+    connection_id = clearskies.columns.String()
+    input_output = clearskies.di.inject.InputOutput()
+
+    def send(self, message):
+        if not self:
+            raise ValueError("Cannot send message to non-existent connection.")
+        if not self.connection_id:
+            raise ValueError(
+                f"Hmmm... I couldn't find the connection id for the {self.__class__.__name__}.  I'm picky about id column names.  Can you please make sure I have a column called connection_id and that it contains the connection id?"
+            )
+
+        domain = self.input_output.context_specifics()["domain"]
+        stage = self.input_output.context_specifics()["stage"]
+        # only include the stage if we're using the default AWS domain - not with a custom domain
+        if ".amazonaws.com" in domain:
+            endpoint_url = f"https://{domain}/{stage}"
+        else:
+            endpoint_url = f"https://{domain}"
+        api_gateway = self.boto3.client("apigatewaymanagementapi", endpoint_url=endpoint_url)
+
+        bytes_message = json.dumps(message).encode("utf-8")
+        try:
+            response = api_gateway.post_to_connection(Data=bytes_message, ConnectionId=self.connection_id)
+        except api_gateway.exceptions.GoneException:
+            self.delete()
+        return response

clearskies_aws/secrets/__init__.py

@@ -1,7 +1,13 @@
-from clearskies import BindingConfig
-from .parameter_store import ParameterStore
-from .secrets_manager import SecretsManager
-from .akeyless_with_ssm_cache import AkeylessWithSsmCache
-from . import additional_configs
-def akeyless_with_ssm_cache(*args, **kwargs):
-    return BindingConfig(AkeylessWithSsmCache, *args, **kwargs)
+from clearskies_aws.secrets import additional_configs
+from clearskies_aws.secrets.akeyless_with_ssm_cache import AkeylessWithSsmCache
+from clearskies_aws.secrets.parameter_store import ParameterStore
+from clearskies_aws.secrets.secrets import Secrets
+from clearskies_aws.secrets.secrets_manager import SecretsManager
+
+__all__ = [
+    "Secrets",
+    "ParameterStore",
+    "SecretsManager",
+    "AkeylessWithSsmCache",
+    "additional_configs",
+]

clearskies_aws/secrets/additional_configs/__init__.py

@@ -1,7 +1,9 @@
-from .mysql_connection_dynamic_producer_via_ssh_cert_bastion import MySQLConnectionDynamicProducerViaSSHCertBastion
-from .mysql_connection_dynamic_producer_via_ssm_bastion import MySQLConnectionDynamicProducerViaSSMBastion
 from .iam_db_auth import IAMDBAuth
 from .iam_db_auth_with_ssm import IAMDBAuthWithSSM
+from .mysql_connection_dynamic_producer_via_ssh_cert_bastion import MySQLConnectionDynamicProducerViaSSHCertBastion
+from .mysql_connection_dynamic_producer_via_ssm_bastion import MySQLConnectionDynamicProducerViaSSMBastion
+
+
 def mysql_connection_dynamic_producer_via_ssh_cert_bastion(
     producer_name=None,
     bastion_host=None,
@@ -26,6 +28,8 @@ def mysql_connection_dynamic_producer_via_ssh_cert_bastion(
         database_name=database_name,
         local_proxy_port=local_proxy_port,
     )
+
+
 def mysql_connection_dynamic_producer_via_ssm_bastion(
     producer_name=None,
     bastion_instance_id=None,
@@ -48,7 +52,11 @@ def mysql_connection_dynamic_producer_via_ssm_bastion(
         database_name=database_name,
         local_proxy_port=local_proxy_port,
     )
+
+
 def iam_db_auth():
     return IAMDBAuth()
+
+
 def iam_db_auth_with_ssm():
     return IAMDBAuthWithSSM()

clearskies_aws/secrets/additional_configs/iam_db_auth.py

@@ -1,29 +1,39 @@
-import clearskies
+from __future__ import annotations
+
 import os
+
+import clearskies
+
+
 class IAMDBAuth(clearskies.di.AdditionalConfig):
     def provide_boto3(self):
         import boto3
+
         return boto3
 
     def provide_connection_details(self, environment, boto3):
         """
-        I really need to make these configurable - both the values themselves and the environment
-        variables that things get pulled from.
+        Make configuration values and environment variables customizable.
+
+        Allows both the values and the environment variable names to be set for flexible configuration.
+
+        Returns:
+            dict: Connection details for IAM DB authentication.
         """
-        endpoint = environment.get('db_endpoint')
-        username = environment.get('db_username')
-        database = environment.get('db_database')
-        region = environment.get('AWS_REGION')
-        ssl_ca_bundle_name = environment.get('ssl_ca_bundle_filename')
-        os.environ['LIBMYSQL_ENABLE_CLEARTEXT_PLUGIN'] = '1'
+        endpoint = environment.get("db_endpoint")
+        username = environment.get("db_username")
+        database = environment.get("db_database")
+        region = environment.get("AWS_REGION")
+        ssl_ca_bundle_name = environment.get("ssl_ca_bundle_filename")
+        os.environ["LIBMYSQL_ENABLE_CLEARTEXT_PLUGIN"] = "1"
 
-        rds_api = boto3.Session().client('rds')
-        rds_token = rds_api.generate_db_auth_token(DBHostname=endpoint, Port='3306', DBUsername=username, Region=region)
+        rds_api = boto3.Session().client("rds")
+        rds_token = rds_api.generate_db_auth_token(DBHostname=endpoint, Port="3306", DBUsername=username, Region=region)
 
         return {
-            'username': username,
-            'password': rds_token,
-            'host': endpoint,
-            'database': database,
-            'ssl_ca': ssl_ca_bundle_name,
+            "username": username,
+            "password": rds_token,
+            "host": endpoint,
+            "database": database,
+            "ssl_ca": ssl_ca_bundle_name,
         }

clearskies_aws/secrets/additional_configs/iam_db_auth_with_ssm.py

@@ -1,57 +1,61 @@
+from __future__ import annotations
+
 import time
+
 import clearskies
+
+
 class IAMDBAuthWithSSM(clearskies.di.AdditionalConfig):
     def provide_subprocess(self):
         import subprocess
+
         return subprocess
 
     def provide_socket(self):
         import socket
+
         return socket
 
     def provide_connection_details(self, environment, subprocess, socket, boto3):
         local_port = self.open_tunnel(environment, subprocess, socket, boto3)
 
         return {
-            'host': '127.0.0.1',
-            'database': environment.get('db_database'),
-            'username': environment.get('db_username'),
-            'password': self.get_password(environment, boto3),
-            'ssl_ca': 'rds-cert-bundle.pem',
-            'port': local_port,
+            "host": "127.0.0.1",
+            "database": environment.get("db_database"),
+            "username": environment.get("db_username"),
+            "password": self.get_password(environment, boto3),
+            "ssl_ca": "rds-cert-bundle.pem",
+            "port": local_port,
         }
 
     def get_password(self, environment, boto3):
-        endpoint = environment.get('db_endpoint')
-        username = environment.get('db_username')
-        region = environment.get('db_region')
+        endpoint = environment.get("db_endpoint")
+        username = environment.get("db_username")
+        region = environment.get("db_region")
 
-        rds_api = boto3.Session().client('rds', region_name=region)
-        return rds_api.generate_db_auth_token(DBHostname=endpoint, Port='3306', DBUsername=username, Region=region)
+        rds_api = boto3.Session().client("rds", region_name=region)
+        return rds_api.generate_db_auth_token(DBHostname=endpoint, Port="3306", DBUsername=username, Region=region)
 
     def open_tunnel(self, environment, subprocess, socket, boto3):
-        endpoint = environment.get('db_endpoint')
-        region = environment.get('db_region')
-        instance_name = environment.get('instance_name')
-        local_proxy_port = int(environment.get('local_proxy_port', '9000'))
+        endpoint = environment.get("db_endpoint")
+        region = environment.get("db_region")
+        instance_name = environment.get("instance_name")
+        local_proxy_port = int(environment.get("local_proxy_port", "9000"))
 
-        ec2_api = boto3.client('ec2', region_name=region)
+        ec2_api = boto3.client("ec2", region_name=region)
         running_instances = ec2_api.describe_instances(
-            Filters=[{
-                'Name': 'tag:Name',
-                'Values': [instance_name]
-            }, {
-                'Name': 'instance-state-name',
-                'Values': ['running']
-            }],
+            Filters=[
+                {"Name": "tag:Name", "Values": [instance_name]},
+                {"Name": "instance-state-name", "Values": ["running"]},
+            ],
         )
         instance_ids = []
-        for reservation in running_instances['Reservations']:
-            for instance in reservation['Instances']:
-                instance_ids.append(instance['InstanceId'])
+        for reservation in running_instances["Reservations"]:
+            for instance in reservation["Instances"]:
+                instance_ids.append(instance["InstanceId"])
 
         if len(instance_ids) == 0:
-            raise ValueError('Failed to launch SSM tunnel! Cannot find bastion!')
+            raise ValueError("Failed to launch SSM tunnel! Cannot find bastion!")
 
         instance_id = instance_ids.pop()
         self._connect_to_bastion(local_proxy_port, instance_id, endpoint, subprocess, socket)
@@ -59,21 +63,21 @@ class IAMDBAuthWithSSM(clearskies.di.AdditionalConfig):
 
     def _connect_to_bastion(self, local_proxy_port, instance_id, endpoint, subprocess, socket):
         sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        result = sock.connect_ex(('127.0.0.1', local_proxy_port))
+        result = sock.connect_ex(("127.0.0.1", local_proxy_port))
         if result == 0:
             sock.close()
             return
 
         tunnel_command = [
-            'aws',
-            '--region',
-            'us-east-1',
-            'ssm',
-            'start-session',
-            '--target',
-            '{}'.format(instance_id),
-            '--document-name',
-            'AWS-StartPortForwardingSessionToRemoteHost',
+            "aws",
+            "--region",
+            "us-east-1",
+            "ssm",
+            "start-session",
+            "--target",
+            "{}".format(instance_id),
+            "--document-name",
+            "AWS-StartPortForwardingSessionToRemoteHost",
             '--parameters={{"host":["{}"], "portNumber":["3306"],"localPortNumber":["{}"]}}'.format(
                 endpoint, local_proxy_port
             ),
@@ -86,7 +90,7 @@ class IAMDBAuthWithSSM(clearskies.di.AdditionalConfig):
             attempts += 1
             time.sleep(0.5)
             sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-            result = sock.connect_ex(('127.0.0.1', local_proxy_port))
+            result = sock.connect_ex(("127.0.0.1", local_proxy_port))
             if result == 0:
                 return
-        raise ValueError('Failed to launch SSM tunnel with command: ' + ' '.join(tunnel_command))
+        raise ValueError("Failed to launch SSM tunnel with command: " + " ".join(tunnel_command))

clearskies_aws/secrets/additional_configs/mysql_connection_dynamic_producer_via_ssh_cert_bastion.py

@@ -1,9 +1,14 @@
-from clearskies.secrets.additional_configs import MySQLConnectionDynamicProducerViaSSHCertBastion as Base
-from pathlib import Path
+from __future__ import annotations
+
+import os
 import socket
 import subprocess
-import os
 import time
+from pathlib import Path
+
+from clearskies.secrets.additional_configs import MySQLConnectionDynamicProducerViaSSHCertBastion as Base
+
+
 class MySQLConnectionDynamicProducerViaSSHCertBastion(Base):
     _config = None
     _boto3 = None
@@ -19,20 +24,20 @@ class MySQLConnectionDynamicProducerViaSSHCertBastion(Base):
         local_proxy_port=None,
         cert_issuer_name=None,
         database_host=None,
-        database_name=None
+        database_name=None,
     ):
         # not using kwargs because I want the argument list to be explicit
         self.config = {
-            'producer_name': producer_name,
-            'bastion_host': bastion_host,
-            'bastion_region': bastion_region,
-            'bastion_name': bastion_name,
-            'bastion_username': bastion_username,
-            'public_key_file_path': public_key_file_path,
-            'local_proxy_port': local_proxy_port,
-            'cert_issuer_name': cert_issuer_name,
-            'database_host': database_host,
-            'database_name': database_name,
+            "producer_name": producer_name,
+            "bastion_host": bastion_host,
+            "bastion_region": bastion_region,
+            "bastion_name": bastion_name,
+            "bastion_username": bastion_username,
+            "public_key_file_path": public_key_file_path,
+            "local_proxy_port": local_proxy_port,
+            "cert_issuer_name": cert_issuer_name,
+            "database_host": database_host,
+            "database_name": database_name,
         }
 
     def provide_connection_details(self, environment, secrets, boto3):
@@ -40,42 +45,36 @@ class MySQLConnectionDynamicProducerViaSSHCertBastion(Base):
         return super().provide_connection_details(environment, secrets)
 
     def _get_bastion_host(self, environment):
-        bastion_host = self._fetch_config(environment, 'bastion_host', 'akeyless_mysql_bastion_host', default='')
-        bastion_name = self._fetch_config(environment, 'bastion_name', 'akeyless_mysql_bastion_name', default='')
+        bastion_host = self._fetch_config(environment, "bastion_host", "akeyless_mysql_bastion_host", default="")
+        bastion_name = self._fetch_config(environment, "bastion_name", "akeyless_mysql_bastion_name", default="")
         if bastion_host:
             return bastion_host
         if bastion_name:
-            bastion_region = self._fetch_config(environment, 'bastion_region', 'akeyless_mysql_bastion_region')
+            bastion_region = self._fetch_config(environment, "bastion_region", "akeyless_mysql_bastion_region")
             return self._public_ip_from_name(bastion_name, bastion_region)
         raise ValueError(
             f"I was asked to connect to a database via an AKeyless dynamic producer through an SSH bastion with certificate auth, but I'm missing some configuration. I need either the bastion host or the name of the instance in AWS.  These can be set in the call to `clearskies.backends.akeyless_aws.mysql_connection_dynamic_producer_via_ssh_cert_bastion()` by providing the 'bastion_host' or 'bastion_name' argument, or by setting an environment variable named 'akeyless_mysql_bastion_host' or 'akeyless_mysql_bastion_name'."
         )
 
     def _public_ip_from_name(self, bastion_name, bastion_region):
-        ec2 = self._boto3.client('ec2', region_name=bastion_region)
+        ec2 = self._boto3.client("ec2", region_name=bastion_region)
         response = ec2.describe_instances(
             Filters=[
-                {
-                    'Name': 'tag:Name',
-                    'Values': [bastion_name]
-                },
-                {
-                    'Name': 'instance-state-name',
-                    'Values': ['running']
-                },
+                {"Name": "tag:Name", "Values": [bastion_name]},
+                {"Name": "instance-state-name", "Values": ["running"]},
             ],
         )
-        if not response.get('Reservations'):
+        if not response.get("Reservations"):
             raise ValueError(
                 f"Could not find a running instance with the designated bastion name, '{bastion_name}' in region '{bastion_region}'"
             )
-        if not response.get('Reservations')[0].get('Instances'):
+        if not response.get("Reservations")[0].get("Instances"):
             raise ValueError(
                 f"Could not find a running instance with the designated bastion name, '{bastion_name}' in region '{bastion_region}'"
             )
-        instance = response.get('Reservations')[0].get('Instances')[0]
-        if not instance.get('PublicIpAddress'):
+        instance = response.get("Reservations")[0].get("Instances")[0]
+        if not instance.get("PublicIpAddress"):
             raise ValueError(
                 f"I found the bastion instance with a name of '{bastion_name}' in region '{bastion_region}', but it doesn't have a public IP address"
             )
-        return instance.get('PublicIpAddress')
+        return instance.get("PublicIpAddress")