piermesh/src/Sponge/base.py

import msgpack
import traceback

from Packets.Message import Message
import Sponge.Protocols.bubble
import Sponge.Protocols.map
import Sponge.Protocols.catch
import Sponge.Protocols.cryptography


class Filter:
    """
    Packet filtering orchestration

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

    cLog
        Reference to `run.Node.cLog` for logging

    cache: Daisy.Cache.Cache
        Reference to our Daisy Cache instance

    completed: list
        List of completed messages IDs

    todo
        Reference to list of actions to do in the Node

    onodeID
        PierMesh node ID
    """

    def __init__(self, cache, onodeID, todo, cLog):
        self.cLog = cLog
        self.cache = cache
        """
        Messages is temporary storage for unfinished messages
        """
        self.messages = {}
        self.completed = []
        self.todo = todo
        self.onodeID = onodeID

    def mCheck(self, payload: bytes):
        """
        Check if payload bytes are msgpack encoded, otherwise skip
        """
        try:
            msgpack.loads(payload)
            return True
        except Exception as e:
            self.cLog(20, "Not msgpack encoded, skipping")
            return False

    def selfCheck(self, packet):
        """
        Check if this is a self packet, if so skip
        """
        if packet["fromId"] == packet["toId"]:
            self.cLog(20, "Self packet, ignored")
            return False
        else:
            return True

    async def protoMap(self, protocolID: int):
        """
        Get protocol from protocol ID using the mlookup table
        """
        protocolID = str(protocolID).zfill(6)
        return self.cache.get("mlookup").get()[protocolID]

    async def protoRoute(self, completeMessage: dict):
        """
        Route message to proper protocol handler
        """
        m = completeMessage
        """
        Shorthand reference completeMessage for ease
        """
        sender = m["sender"]
        senderDisplay = m["senderDisplayName"]
        recipient = m["recipient"]
        recipientNode = m["recipientNode"]
        # TODO: Fix packets to use class
        protocol = await self.protoMap(m["packetsClass"])
        self.cLog(20, "Protocol: " + protocol)
        if protocol == "bubble":
            await Sponge.Protocols.bubble.filter(
                m, recipient, recipientNode, self.onodeID, self.todo
            )
        elif protocol == "map":
            await Sponge.Protocols.map.filter(m, self.todo)
        elif protocol == "catch":
            await Sponge.Protocols.catch.filter(m, recipient, recipientNode, self.todo)
        elif protocol == "cryptography":
            await Sponge.Protocols.cryptography.filter(
                completeMessage, recipientNode, self.todo
            )
        else:
            self.cLog(30, "Cant route, no protocol")

    async def sieve(self, packet):
        """
        Base filtering logic, takes a single MeshTastic packet
        """
        p = packet["decoded"]["payload"]
        if self.selfCheck(packet) and self.mCheck(p):
            try:
                p = msgpack.loads(p)
                self.cLog(20, p)
                packetsID = p["packetsID"]
                packetsClass = p["packetsClass"]
                if packetsID == 0:
                    self.cLog(20, "Single packet")
                    # Do sp handling
                    pass
                if packetsID in self.completed:
                    raise ValueError("Message already completed")
                if not (packetsID in self.messages):
                    self.messages[packetsID] = {
                        "packetCount": p["packetCount"],
                        "data": [],
                        "finished": False,
                        "dataOrder": [],
                    }
                if "wantFullResponse" in p.keys():
                    for k in p.keys():
                        if k != "data":
                            self.messages[packetsID][k] = p[k]
                elif not p["packetNumber"] in self.messages[packetsID]["dataOrder"]:
                    self.messages[packetsID]["data"].append(p["data"])
                    self.messages[packetsID]["dataOrder"].append(p["packetNumber"])
                if (len(self.messages[packetsID]["data"])) >= (
                    self.messages[packetsID]["packetCount"] - 1
                ) and ("wantFullResponse" in self.messages[packetsID].keys()):
                    self.cLog(20, "Finished receiving for message " + str(packetsID))
                    self.messages[packetsID]["finished"] = True
                    if self.messages[packetsID]["wantFullResponse"] != False:
                        # TO DO: implement loop
                        # responseLoop(packets_id)
                        pass
                    completeMessage = self.messages[packetsID]
                    completeMessage["data"] = Message.reassemble(None, completeMessage)
                    del self.messages[packetsID]
                    self.completed.append(packetsID)
                    self.cLog(20, "Assembly completed, routing")
                    await self.protoRoute(completeMessage)
            except Exception:
                self.cLog(30, traceback.format_exc())