[PRM-495] Moved location of the awsmesh.entrypoint module

Author: oliverbeumkes-nhs

services/mesh-forwarder/__init__.py

@@ -0,0 +1,53 @@
+import logging
+import sys
+from dataclasses import MISSING, Field, dataclass, fields
+from distutils.util import strtobool
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+
+def _parse_field_from_env_var(name: str, value: str, field: Field):
+    if field.type == Optional[bool]:
+        try:
+            return strtobool(value)
+        except ValueError:
+            logger.warning(
+                f"Invalid value '{value}' for {name}, ignoring and using default: {field.default}"
+            )
+            return field.default
+    return value
+
+
+def _read_env(field: Field, env_vars):
+    env_var_name = field.name.upper()
+    if env_var_name in env_vars:
+        return _parse_field_from_env_var(env_var_name, env_vars[env_var_name], field)
+    elif field.default != MISSING:
+        return field.default
+    else:
+        logger.error(f"Expected environment variable {env_var_name} was not set, exiting...")
+        sys.exit(1)
+
+
+@dataclass
+class ForwarderConfig:
+    mesh_url: str
+    mesh_mailbox_ssm_param_name: str
+    mesh_password_ssm_param_name: str
+    mesh_shared_key_ssm_param_name: str
+    mesh_client_cert_ssm_param_name: str
+    mesh_client_key_ssm_param_name: str
+    mesh_ca_cert_ssm_param_name: str
+    poll_frequency: str
+    forwarder_home: str
+    message_destination: str = "s3"
+    s3_bucket_name: Optional[str] = None
+    sns_topic_arn: Optional[str] = None
+    endpoint_url: Optional[str] = None
+    ssm_endpoint_url: Optional[str] = None
+    disable_message_header_validation: Optional[bool] = False
+
+    @classmethod
+    def from_environment_variables(cls, env_vars):
+        return cls(**{field.name: _read_env(field, env_vars) for field in fields(cls)})

services/mesh-forwarder/awsmesh/__init__.py

@@ -0,0 +1,89 @@
+import logging
+from os import environ
+from os.path import join
+from signal import SIGINT, SIGTERM, signal
+
+import boto3
+import urllib3
+
+from awsmesh.config import ForwarderConfig
+from awsmesh.forwarder_service import MeshConfig, build_forwarder_service
+from awsmesh.logging import JsonFormatter
+from awsmesh.message_destination_resolver import MessageDestinationConfig
+from awsmesh.secrets import SsmSecretManager
+
+urllib3.disable_warnings(urllib3.exceptions.SubjectAltNameWarning)
+
+
+def build_mesh_config_from_ssm(ssm, config) -> MeshConfig:
+    mesh_client_cert_path = join(config.forwarder_home, "client_cert.pem")
+    mesh_client_key_path = join(config.forwarder_home, "client_key.pem")
+    mesh_ca_cert_path = join(config.forwarder_home, "ca_cert.pem")
+
+    secret_manager = SsmSecretManager(ssm)
+
+    secret_manager.download_secret(config.mesh_client_cert_ssm_param_name, mesh_client_cert_path)
+    secret_manager.download_secret(config.mesh_client_key_ssm_param_name, mesh_client_key_path)
+    secret_manager.download_secret(config.mesh_ca_cert_ssm_param_name, mesh_ca_cert_path)
+
+    mesh_mailbox = secret_manager.get_secret(config.mesh_mailbox_ssm_param_name)
+    mesh_password = secret_manager.get_secret(config.mesh_password_ssm_param_name)
+    mesh_shared_key = secret_manager.get_secret(config.mesh_shared_key_ssm_param_name)
+
+    return MeshConfig(
+        url=config.mesh_url,
+        mailbox=mesh_mailbox,
+        password=mesh_password,
+        shared_key=bytes(mesh_shared_key, "utf-8"),
+        client_cert_path=mesh_client_cert_path,
+        client_key_path=mesh_client_key_path,
+        ca_cert_path=mesh_ca_cert_path,
+    )
+
+
+def build_message_destination_config(config) -> MessageDestinationConfig:
+    return MessageDestinationConfig(
+        message_destination=config.message_destination,
+        s3_bucket_name=config.s3_bucket_name,
+        endpoint_url=config.endpoint_url,
+        sns_topic_arn=config.sns_topic_arn,
+    )
+
+
+def build_forwarder_from_environment_variables(env_vars=environ):
+    config = ForwarderConfig.from_environment_variables(env_vars)
+    ssm = boto3.client("ssm", endpoint_url=config.ssm_endpoint_url)
+
+    return build_forwarder_service(
+        mesh_config=build_mesh_config_from_ssm(ssm, config),
+        message_destination_config=build_message_destination_config(config),
+        poll_frequency_sec=int(config.poll_frequency),
+        disable_message_header_validation=config.disable_message_header_validation,
+    )
+
+
+def setup_logger():
+    logger = logging.getLogger()
+    logger.setLevel(logging.INFO)
+    formatter = JsonFormatter()
+    handler = logging.StreamHandler()
+    handler.setFormatter(formatter)
+    logger.addHandler(handler)
+
+
+def main():
+    setup_logger()
+
+    forwarder_service = build_forwarder_from_environment_variables()
+
+    def handle_sigterm(signum, frame):
+        forwarder_service.stop()
+
+    signal(SIGINT, handle_sigterm)
+    signal(SIGTERM, handle_sigterm)
+
+    forwarder_service.start()
+
+
+if __name__ == "__main__":
+    main()

services/mesh-forwarder/awsmesh/config.py

@@ -0,0 +1,85 @@
+import logging
+
+from awsmesh.mesh import InvalidMeshHeader, MeshClientNetworkError, MeshInbox, MissingMeshHeader
+from awsmesh.monitoring.probe import LoggingProbe
+from awsmesh.uploader import MessageUploader, UploaderError
+
+logger = logging.getLogger(__name__)
+
+
+class RetryableException(Exception):
+    pass
+
+
+class MeshToAwsForwarder:
+    def __init__(
+        self,
+        inbox: MeshInbox,
+        uploader: MessageUploader,
+        probe: LoggingProbe,
+        disable_message_header_validation,
+    ):
+        self._inbox = inbox
+        self._uploader = uploader
+        self._probe = probe
+        self._disable_message_header_validation = disable_message_header_validation
+
+    def forward_messages(self):
+        retryable_message_exceptions = []
+        for message_id in self._poll_message_ids():
+            try:
+                self._process_message(message_id)
+            except RetryableException as e:
+                retryable_message_exceptions.append(e)
+
+        if len(retryable_message_exceptions) > 0:
+            logger.info(
+                f"Raising single retryable exception, actually caught {len(retryable_message_exceptions)} message exception(s)"
+            )
+            raise retryable_message_exceptions[0]
+
+    def is_mailbox_empty(self):
+        count_message_event = self._probe.new_count_messages_event()
+        try:
+            message_count = self._inbox.count_messages()
+            count_message_event.record_message_count(message_count)
+            return message_count <= 0
+        except MeshClientNetworkError as e:
+            count_message_event.record_mesh_client_network_error(e)
+            raise RetryableException()
+        finally:
+            count_message_event.finish()
+
+    def _poll_message_ids(self):
+        poll_inbox_event = self._probe.new_poll_inbox_event()
+        try:
+            messages = self._inbox.list_message_ids()
+            poll_inbox_event.record_message_batch_count(len(messages))
+            return messages
+        except MeshClientNetworkError as e:
+            poll_inbox_event.record_mesh_client_network_error(e)
+            raise RetryableException()
+        finally:
+            poll_inbox_event.finish()
+
+    # flake8: noqa: C901
+    def _process_message(self, message_id):
+        forward_message_event = self._probe.new_forward_message_event()
+        try:
+            message = self._inbox.retrieve_message(message_id)
+            forward_message_event.record_message_metadata(message)
+            if not self._disable_message_header_validation:
+                message.validate()
+            self._uploader.upload(message, forward_message_event)
+            message.acknowledge()
+        except MissingMeshHeader as e:
+            forward_message_event.record_missing_mesh_header(e)
+        except InvalidMeshHeader as e:
+            forward_message_event.record_invalid_mesh_header(e)
+        except UploaderError as e:
+            forward_message_event.record_uploader_error(e)
+        except MeshClientNetworkError as e:
+            forward_message_event.record_mesh_client_network_error(e)
+            raise RetryableException()
+        finally:
+            forward_message_event.finish()

services/mesh-forwarder/awsmesh/entrypoint.py

@@ -0,0 +1,75 @@
+import logging
+from dataclasses import dataclass
+from threading import Event
+from typing import Optional
+
+import mesh_client
+
+from awsmesh.forwarder import MeshToAwsForwarder, RetryableException
+from awsmesh.mesh import MeshInbox
+from awsmesh.message_destination_resolver import MessageDestinationConfig, resolve_message_uploader
+from awsmesh.monitoring.probe import LoggingProbe
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class MeshConfig:
+    url: str
+    mailbox: str
+    password: str
+    shared_key: bytes
+    client_cert_path: str
+    client_key_path: str
+    ca_cert_path: str
+
+
+class MeshToAwsForwarderService:
+    def __init__(
+        self,
+        forwarder: MeshToAwsForwarder,
+        poll_frequency_sec: int,
+        exit_event: Optional[Event] = None,
+    ):
+        self._forwarder = forwarder
+        self._exit_event = exit_event or Event()
+        self._poll_frequency_sec = poll_frequency_sec
+
+    def start(self):
+        logger.info("Started forwarder service")
+        while not self._exit_event.is_set():
+            try:
+                self._forwarder.forward_messages()
+
+                if self._forwarder.is_mailbox_empty():
+                    self._exit_event.wait(self._poll_frequency_sec)
+            except RetryableException:
+                self._exit_event.wait(self._poll_frequency_sec)
+        logger.info("Exiting forwarder service")
+
+    def stop(self):
+        logger.info("Received request to stop")
+        self._exit_event.set()
+
+
+def build_forwarder_service(
+    mesh_config: MeshConfig,
+    message_destination_config: MessageDestinationConfig,
+    poll_frequency_sec: int,
+    disable_message_header_validation: bool,
+) -> MeshToAwsForwarderService:
+    uploader = resolve_message_uploader(message_destination_config)
+
+    mesh = mesh_client.MeshClient(
+        mesh_config.url,
+        mesh_config.mailbox,
+        mesh_config.password,
+        shared_key=mesh_config.shared_key,
+        cert=(mesh_config.client_cert_path, mesh_config.client_key_path),
+        verify=mesh_config.ca_cert_path,
+    )
+    inbox = MeshInbox(mesh)
+    forwarder = MeshToAwsForwarder(
+        inbox, uploader, LoggingProbe(), disable_message_header_validation
+    )
+    return MeshToAwsForwarderService(forwarder, poll_frequency_sec)

services/mesh-forwarder/awsmesh/forwarder.py

@@ -0,0 +1,26 @@
+import json
+from datetime import datetime
+from logging import LogRecord, makeLogRecord
+
+DEFAULT_LOG_RECORD_ATTRS = vars(makeLogRecord({})).keys()
+
+
+def _convert_timestamp_to_iso(timestamp: float) -> str:
+    return datetime.utcfromtimestamp(timestamp).isoformat()
+
+
+class JsonFormatter:
+    def format(self, record: LogRecord) -> str:
+        base = {
+            "level": record.levelname,
+            "module": record.module,
+            "message": record.msg,
+            "time": _convert_timestamp_to_iso(record.created),
+        }
+        extra = {
+            attr: getattr(record, attr)
+            for attr in vars(record)
+            if attr not in DEFAULT_LOG_RECORD_ATTRS
+        }
+
+        return json.dumps({**base, **extra})