import base64 import re from io import BytesIO from datetime import datetime, timezone from typing import Dict, List from amqp.model.model import AMQMessage from amqp.model.snowflake_id import SnowflakeId from amqp.router.serializer import map_with_list_as_string, map_as_string, parse_map_list_string, parse_map_string from amqp.router.utils import sanitize_as_rabbitmq_name class AMQMessageFactory: SNOWFLAKE_ID_BITS = 80 SEQUENCE_BITS = 12 MACHINE_ID_BITS = 24 TIMESTAMP_BITS = SNOWFLAKE_ID_BITS - SEQUENCE_BITS - MACHINE_ID_BITS MACHINE_ID_MASK = (1 << MACHINE_ID_BITS) - 1 SEQUENCE_MASK = (1 << SEQUENCE_BITS) - 1 MAGIC_BYTE_HTTP_REQUEST = "A" _instance = None snowflake_epoch = datetime(2024, 10, 1, 0, 0, tzinfo=timezone.utc).timestamp() def __init__(self, machine_id: int): self.machine_id = (machine_id & self.MACHINE_ID_MASK) << self.SEQUENCE_BITS self.snowflake_sequence = 0 self.snowflake_last_timestamp = 0 @classmethod def get_instance(cls, machine_id: int): """ return reusable factory instance :param machine_id: a generator ID, should be unique for each component to prevent duplicate IDs :return: factory instance """ if cls._instance is None: cls._instance = AMQMessageFactory(machine_id) return cls._instance def generate_snowflake_id(self) -> SnowflakeId: """ Generate an unique ID following the SnowflakeID structure, see https://en.wikipedia.org/wiki/Snowflake_ID :return: 80-bit wide ID """ timestamp = int(datetime.now(timezone.utc).timestamp() * 1000) if timestamp < self.snowflake_last_timestamp: raise RuntimeError("Clock moved backwards. Refusing to generate id") if timestamp == self.snowflake_last_timestamp: self.snowflake_sequence = (self.snowflake_sequence + 1) & self.SEQUENCE_MASK if self.snowflake_sequence == 0: while timestamp == self.snowflake_last_timestamp: timestamp = int(datetime.now(timezone.utc).timestamp() * 1000) self.snowflake_last_timestamp = timestamp else: self.snowflake_sequence = 0 self.snowflake_last_timestamp = timestamp time = (timestamp - 1000 * int(AMQMessageFactory.snowflake_epoch)) return SnowflakeId( time >> (64 - self.SEQUENCE_BITS - self.MACHINE_ID_BITS), ((time << (self.SEQUENCE_BITS + self.MACHINE_ID_BITS)) | self.machine_id | self.snowflake_sequence) & 0xFFFFFFFFFFFFFFFF ) def create_request_message( self, method: str, domain: str, path: str, headers: Dict[str, List[str]], message, ) -> AMQMessage: """ Encapsulate / hide the system level inputs, create the AMQ message from business relevant input :param method: a method requested (e.g. HTTP method like GET, POST in case of REST request) :param domain: service host name :param path: path at the service :param headers: HTTP headers :param message: payload :return: AMQMessage object """ serializable_headers = headers.copy() if headers else {} payload = base64.b64encode(message.encode('utf-8') if isinstance(message, str) else message) trace_info = {} return AMQMessage( self.MAGIC_BYTE_HTTP_REQUEST, self.generate_snowflake_id(), method, domain, path, serializable_headers, trace_info, payload, ) @staticmethod def amq_message_routing_key(message: AMQMessage) -> str: """ Constructs the message specific routing key. :param message: AMQMessage input :return: routing key (as string) compliant to RabbitMQ allowed character set. """ return sanitize_as_rabbitmq_name(f"{message.domain}.{message.path}") @staticmethod def get_timestamp_from_id(_id: SnowflakeId) -> int: """ Helper method - retrieve the timestamp from the ID :param _id: SnowflakeID (includes encoded timestamp) :return: the message timestamp """ bits = _id.get_bits() low = bits[0] >> (AMQMessageFactory.SEQUENCE_BITS + AMQMessageFactory.MACHINE_ID_BITS) high = bits[1] << (64 - AMQMessageFactory.SEQUENCE_BITS - AMQMessageFactory.MACHINE_ID_BITS) return (high | low) + AMQMessageFactory.snowflake_epoch * 1000 @staticmethod def to_string(message: AMQMessage) -> str: """ Human-readable representation. :param message: AMQMessage input :return: as string """ return ( f"AMQMessage{{" f"id={message.id}, " f"method='{message.method}', " f"domain='{message.domain}', " f"path='{message.path}', " f"headers={{{map_with_list_as_string(message.headers)}}}, " f"traceInfo={{{map_as_string(message.trace_info)}}}, " f"base64body={message.base64_body.decode()}" f"}}" ) @staticmethod def serialize(message: AMQMessage) -> bytes: """ Converts AMQMessage to byte array :param message: message to serialize :return: messages as bytes """ id_bytes = str(message.id).encode() prolog = ( f"|m={message.method}|d={message.domain}|p={message.path}|h={{{map_with_list_as_string(message.headers)}}}|t={{{map_as_string(message.trace_info)}}}|b=" ).encode('utf-8') header_length = 1 + len(id_bytes) + len(prolog) prolog_len = header_length.to_bytes(2, byteorder='big') msg_length = 2 + header_length + len(message.base64_body) with BytesIO(bytearray(msg_length)) as baos: baos.write(prolog_len) baos.write(AMQMessageFactory.MAGIC_BYTE_HTTP_REQUEST.encode('utf-8')) baos.write(id_bytes) baos.write(prolog) baos.write(message.base64_body) return baos.getvalue() @staticmethod def from_bytes(input_bytes: bytes) -> AMQMessage: """ Builds the AMQMessage from its serialized (byte array) form. :param input_bytes: serialized message :return: AMQMessage instance """ prolog_len = (input_bytes[0] << 8) + input_bytes[1] metadata = input_bytes[2:prolog_len + 2].decode() pattern = re.compile( f"{AMQMessageFactory.MAGIC_BYTE_HTTP_REQUEST}" r"([0-9A-Fa-f]+)\|" r"m=([^|]*)\|" r"d=([^|]*)\|" r"p=([^|]*)\|" r"h=(\{[^}]*\})\|" r"t=(\{[^}]*\})\|" r"b=" ) match = pattern.match(metadata) if not match: raise ValueError( f"AMQMessage format does not match expected format. " f"Input string [{metadata}] does not match pattern: {pattern.pattern}" ) id_str = match.group(1) method = match.group(2) domain = match.group(3) path = match.group(4) headers_str = match.group(5) trace_info_str = match.group(6) body_b64 = input_bytes[prolog_len + 2:] headers = parse_map_list_string(headers_str) trace_info = parse_map_string(trace_info_str) return AMQMessage( AMQMessageFactory.MAGIC_BYTE_HTTP_REQUEST, SnowflakeId.from_hex(id_str), method, domain, path, headers, trace_info, body_b64, )