piermesh/src/Packets/Message.py

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

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

logger = logging.getLogger("__main__." + __name__)

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,
        target=True,
        subMessage=False,
        primaryMessage=None,
        pskEncrypt=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
        """
        self.recipientNode = recipientNode
        self.target = target
        self.subMessage = subMessage
        if subMessage:
            self.primaryMessage = primaryMessage
        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)
            if subMessage == False:
                bytesObject, nonce, tag = cryptographyInfo.encrypt(
                    bytesObject, str(recipientNode).zfill(6), isDict=False, pskEncrypt=pskEncrypt
                )
                logger.log(10, bytesObject)
                self.nonce = nonce
                self.tag = tag
            packets = []
            self.packetsID = random.randrange(0, 999999)
            pnum = 1
            # if subMessage:
            dataSize = 80
            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)]
                if subMessage:
                    packets.append(
                        p.Packet(
                            b,
                            self.packetsID,
                            pnum,
                            packetsClass=packetsClass,
                            primaryMessage=primaryMessage,
                        )
                    )
                else:
                    packets.append(
                        p.Packet(b, self.packetsID, pnum, packetsClass=packetsClass)
                    )
                pnum += 1
                tb += b
            if subMessage:
                pass
            else:
                packets.insert(
                    0,
                    h.Header(
                        self.packetsID,
                        pnum,
                        sender,
                        senderDisplayName,
                        sourceNode,
                        recipient,
                        recipientNode,
                        wantFullResponse=wantFullResponse,
                        packetsClass=packetsClass,
                        pAction=pAction,
                        target=target,
                    ),
                )

            self.fullPackets = [p for p in packets]

            if subMessage:
                pnum -= 1
                
            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, subMessage=False, yctx=None, packetCount=None):
        """
        Reassemble packets from a completed message in `Sponge.base`
        """
        data = b""
        sourceNode = None
        # TODO: Fix reassembly for primary
        if subMessage:
            sourceNode = yctx["sourceNode"]["val"]
            for it in range(1, packetCount+1):
                data += completedMessage["data"][completedMessage["dataOrder"].index(it)]
            data = msgpack.loads(lzma.decompress(data))
            logger.log(10, data)
            logger.log(10, completedMessage["data"])
            logger.log(10, completedMessage["dataOrder"])
        else:
            # TODO: Cryptography setup
            packetCount = int(completedMessage.yctx["packetCount"]["val"])
            sourceNode = completedMessage.yctx["sourceNode"]["val"]
            logger.log(10, completedMessage.data)
            for it in range(1, packetCount):
                if it in completedMessage.dataOrder:
                    data += completedMessage.data[completedMessage.dataOrder.index(it)]
            logger.log(10, "pre decrypt")
            logger.log(10, data)
            data = cryptographyInfo.decrypt(
                data, sourceNode, completedMessage.nonce, completedMessage.tag
            )
            # data = msgpack.loads(lzma.decompress(data))
        return data