How to client-side IAM auth an Amazon API Gateway WebSocket connection

APIGW_websocket_IAM_Auth_HowTo

    async def connect():
        # Create a version of the websocket client class that handles AWS sigv4
        # authorization by overriding the 'write_http_request' method with the
        # logic to construct an x-amzn-auth header at the last possible moment.
        def class WebSocketSigv4ClientProtocol(WebSocketClientProtocol):
            def __init__(self, *args, **kwargs) -> None:
                super().__init__(*args, **kwargs)
            def write_http_request(self, path: str, headers) -> None:
                # Intercept the GET that initiates the websocket protocol at the point where
                # all of its 'real' headers have been constructed. Add in the sigv4 header AWS needs.
                credentials = Credentials(
                    os.environ['AWS_ACCESS_KEY_ID'],
                    os.environ['AWS_SECRET_ACCESS_KEY'],
                    os.environ['AWS_SESSION_TOKEN'])
                sigv4 = SigV4Auth(credentials, 'execute-api', os.environ['AWS_REGION'])
                request = AWSRequest(method='GET', url='https://' + natpunch_server)
                sigv4.add_auth(request)
                prepped = request.prepare()
                headers['Authorization'       ] = prepped.headers['Authorization'       ]
                headers['X-Amz-Date'          ] = prepped.headers['X-Amz-Date'          ]
                headers['x-amz-security-token'] = prepped.headers['x-amz-security-token']
                # Run the original code with the added sigv4 auth header now included:
                super().write_http_request(path, headers)

        if (not 'AWS_ACCESS_KEY_ID' in os.environ):
            raise Exception('missing environment variable(s) required for signing',
                            'AWS_ACCESS_KEY_ID not present')
        if (not 'AWS_SECRET_ACCESS_KEY' in os.environ):
            raise Exception('missing environment variable(s) required for signing',
                            'AWS_SECRET_ACCESS_KEY not present')
        if (not 'AWS_SESSION_TOKEN' in os.environ):
            raise Exception('missing environment variable(s) required for signing',
                            'AWS_SESSION_TOKEN not present')
        if (not 'AWS_REGION' in os.environ):
            raise Exception('missing environment variable(s) required for signing',
                            'AWS_REGION not present')

        async with websockets.connect('wss://...', # FILL THIS IN 
                                      create_protocol:WebSocketSigv4ClientProtocol,
                                      extra_headers={'x-api-key':api_key}) as websocket:
            msg_as_string = json.dumps({...}) # FILL THIS IN
            await websocket.send(msg_as_string)
            try:
                return json.loads(await asyncio.wait_for(websocket.recv(), timeout=my_timeout))
            except asyncio.TimeoutError:
                return None
return asyncio.run(connect())