leonet/leonet/frames.py

import typing
import datetime
import dataclasses

import crypto
import message_codec


def create_frame(frame_type, recipient, originator, content, password):
    main_frame = f"{frame_type}#{recipient}#{originator}#{content}"

    checksum = crypto.find_checksum(main_frame)
    auth_code = crypto.find_auth_code(password, main_frame)

    return f"\nLEONET {main_frame} {checksum} {auth_code} TENOEL\n"


@dataclasses.dataclass
class MsgFrame:
    recipient: str
    originator: str
    message: str
    timestamp: dataclasses.InitVar[typing.Optional[str]] = None
    message_id: dataclasses.InitVar[typing.Optional[str]] = None

    def __post_init__(self, message_id, timestamp):
        if timestamp is not None:
            self.timestamp = timestamp
        else:
            self.timestamp = datetime.datetime.utcnow().isoformat()

        if message_id is not None:
            self.message_id = message_id
        else:
            safe_message = message_codec.encode(self.message)
            body = f"{self.timestamp}&{safe_message}"
            self.message_id = crypto.find_checksum(body)

        self.frame_recipient = self.recipient

    def encode(self, password=""):
        safe_message = message_codec.encode(self.message)
        body = f"{self.timestamp}&{safe_message}"
        message_id = crypto.find_checksum(body)
        full_body = f"{body};{message_id}"

        return create_frame(
            "MSG", self.recipient, self.originator, full_body, password
        )

    def gen_ack(self):
        return frames.AckFrame(
            self.recipient,
            self.originator,
            self.message_id,
        )

    def gen_nak(self):
        return frames.NakFrame(
            self.recipient,
            self.originator,
            self.message_id,
        )


@dataclasses.dataclass
class QryFrame:
    recipient: str
    originator: str
    message_id: str

    def __post_init__(self):
        self.frame_recipient = self.recipient

    def encode(self, password=""):
        return create_frame(
            "QRY", self.recipient, self.originator, self.message_id, password
        )

    def gen_ack(self):
        return frames.AckFrame(
            self.recipient,
            self.originator,
            self.message_id,
        )

    def gen_nak(self):
        return frames.NakFrame(
            self.recipient,
            self.originator,
            self.message_id,
        )


@dataclasses.dataclass
class AckFrame:
    recipient: str
    originator: str
    message_id: str

    def __post_init__(self):
        self.frame_recipient = self.originator

    def encode(self, password=""):
        return create_frame(
            "ACK", self.recipient, self.originator, self.message_id, password
        )


@dataclasses.dataclass
class NakFrame:
    recipient: str
    originator: str
    message_id: str

    def __post_init__(self):
        self.frame_recipient = self.originator

    def encode(self, password=""):
        return create_frame(
            "NAK", self.recipient, self.originator, self.message_id, password
        )