piermesh/src/Sponge/base.py

import msgpack
import traceback

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


class Filter:
    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):
        try:
            msgpack.loads(payload)
            return True
        except Exception as e:
            self.cLog(20, "Not msgpack encoded, skipping")
            return False

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

    async def protoMap(self, protocolID):
        protocolID = str(protocolID).zfill(6)
        return self.cache.get("mlookup").get()[protocolID]

    async def protoRoute(self, completeMessage):
        """
        Shorthand reference
        """
        m = completeMessage
        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):
        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
                    # TODO: Sorting
                    completeMessage = self.messages[packetsID]
                    completeMessage["data"] = Packets.reassemble(None, completeMessage)
                    del self.messages[packetsID]
                    self.completed.append(packetsID)
                    self.cLog(20, "Assembly completed, routing")
                    # self.cache.create("messages/" + str(packetsID), cm)
                    await self.protoRoute(completeMessage)
            except Exception as e:
                self.cLog(30, traceback.print_exc())