piermesh/src/Packets/Message.py

import Packets.Packet as p
import Packets.HeaderPacket as h
import lzma
import msgpack
import random
import math

# DO NOT CHANGE DATA SIZE UNLESS YOU KNOW WHAT YOURE DOING


def dict2bytes(cdict: dict):
    return lzma.compress(msgpack.dumps(cdict))


class Message:
    """
    Full message which is composed of `Packets.Packet.Packet`s

    `🔗 Source <https://git.utopic.work/PierMesh/piermesh/src/branch/main/Packets/Message.py>`__

    Attributes
    ----------
    packets: list[Packets.Packet.Packet]
        List of packets making up the Message
    """

    def __init__(
        self,
        bytesObject: bytes,
        sender: int,
        senderDisplayName: int,
        sourceNode,
        recipient: int,
        recipientNode: int,
        cryptographyInfo,
        packetsClass,
        pAction,
        dataSize: int = 128,
        wantFullResponse: bool = False,
    ):
        # TODO: PSK for usage prior to credentials
        """
        Parameters
        ----------
        bytesObject: bytes
            Bytes to split into packets

        sender: int
            6 digit (maximum) node or peer ID

        senderDisplayName: int
            3 digit (maximum) ID for mapping display names to a given user

        recipient: int
            6 digit (maximum) node or peer ID

        recipientNode: int
            6 digit (maximum) node ID to route the packet to

        dataSize: int
            Size to cut the bytesObject into per packet

        wantFullResponse: bool
            Whether to send a response when the message has completed reception (TODO: Kill all retries for associated packets when received)

        packetsClass: int
            Which protocol the packets are using
        """
        if isinstance(bytesObject, list):
            packets = [h.Header(bytesObject[0])]
            for packet in bytesObject:
                packets.append(
                    p.Packet(
                        packet["data"],
                        packetsID=packet["packetsID"],
                        packetNumber=packet["packetNumber"],
                        packetsClass=packetsClass,
                    )
                )
            self.packets = packets
        else:
            # Data passed in by peers should already have been e2ee encrypted by SubtleCrypto
            # Transport encryption
            # bytesObject = lzma.compress(bytesObject, str(recipientNode).zfill(6), isDict=False)
            bytesObject, nonce, tag = cryptographyInfo.encrypt(
                bytesObject, str(recipientNode).zfill(6), isDict=False
            )
            packets = []
            self.packetsID = random.randrange(0, 999999)
            pnum = 1
            blen = math.ceil(len(bytesObject) / dataSize)
            tb = b""
            for it in range(blen):
                if it >= (blen - 1):
                    b = bytesObject[it * dataSize :]
                else:
                    b = bytesObject[it * dataSize : (it * dataSize + dataSize)]
                packets.append(
                    p.Packet(b, self.packetsID, pnum, packetsClass=packetsClass)
                )
                pnum += 1
                tb += b
            packets.insert(
                0,
                h.Header(
                    self.packetsID,
                    pnum,
                    sender,
                    senderDisplayName,
                    sourceNode,
                    recipient,
                    recipientNode,
                    nonce,
                    wantFullResponse=wantFullResponse,
                    packetsClass=packetsClass,
                    pAction=pAction,
                ),
            )
            for it in range(pnum):
                packet = msgpack.loads(packets[it].dump())
                packet["packetCount"] = pnum

                packets[it] = msgpack.dumps(packet)

            self.packets = packets

    def get(self) -> list[p.Packet]:
        """
        Get and return all packets
        """
        return self.packets

    def reassemble(self, completedMessage: dict, cryptographyInfo):
        """
        Reassemble packets from a completed message in `Sponge.base`
        """
        data = b""
        for it in range(1, int(completedMessage["packetCount"])):
            data += completedMessage["data"][completedMessage["dataOrder"].index(it)]
        res = msgpack.loads(
            lzma.decompress(
                cryptographyInfo.decrypt(
                    data, completedMessage["sourceNode"], completedMessage["nonce"]
                )
            )
        )
        return res