piermesh/src/Sponge/Protocols/Yellow.py

from Services.Action import Action
from Config.Context import Context
from Packets.Message import Message

import logging

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


class YCTX(Context):
    def __init__(
        self,
        packetsID,
        packetCount,
        pAction,
        todo,
        cryptographyInfo,
        sourceNode,
        subMessage=False,
        subMessages={},
        submessagesIDs=[],
        eData=None,
    ):
        super().__init__(
            sourceNode=sourceNode,
            packetsID=packetsID,
            packetCount=packetCount,
            pAction=pAction,
            cryptographyInfo=cryptographyInfo,
            subMessages=subMessages,
            submessagesIDs=submessagesIDs,
            eData=eData,
            subMessage=subMessage,
        )
        self.todo = todo


class Yellow:
    # TODO: Submessage completion actions
    pActions = []

    def __init__(
        self,
        yctx: YCTX,
    ):
        self.yctx = yctx
        self.message = None
        self.submessages = yctx["subMessages"]["val"]
        self.submessagesIDs = yctx["submessagesIDs"]["val"]
        self.finishedSubmessages = {}
        self.dataOrder = []
        self.data = []
        self.nonce = None
        self.tag = None
        self.gotHead = False
        self.todo = yctx.todo
        if yctx["eData"]["val"] != None:
            self.dataOrder = yctx["eData"]["val"]["dataOrder"]
            self.data = yctx["eData"]["val"]["data"]

    def checkComplete(self):
        logger.log(30, "In check complete")
        for submessage in self.submessagesIDs:
            submessage = str(submessage)
            if submessage not in self.finishedSubmessages.keys():
                logger.log(30, f"Submessage {submessage} missing")
                return False        
        logger.log(30, "Submessages complete")
        return True

    async def id(self):
        return self.packetsID

    async def processPacket(self, p, subMessage=False, rdoaoc=[]):
        doCheck = True
        if subMessage == False and p["packetNumber"] != False:
            #if not p["packetNumber"] in self.dataOrder:
            if p["packetNumber"] not in self.dataOrder:
                self.data.append(p["data"])
                self.dataOrder.append(p["packetNumber"])
                logger.log(10, "data")
                logger.log(10, self.data)
            else:
                doCheck = False
        elif p["packetNumber"] == False:
            if not self.gotHead:
                self.data.append(False)
                self.dataOrder.append(False)
                self.submessagesIDs = p["submessages"]
                self.gotHead = True
            else:
                doCheck = False
        else:
            if not str(p["packetsID"]) in self.submessages.keys():
                self.submessages[str(p["packetsID"])] = {
                    "data": [],
                    "dataOrder": []
                }
            if p["packetNumber"] not in self.submessages[str(p["packetsID"])]["dataOrder"]:
                self.submessages[str(p["packetsID"])]["data"].append(p["data"])
                self.submessages[str(p["packetsID"])]["dataOrder"].append(
                p["packetNumber"])
            else:
                doCheck = False
        if doCheck:
            if subMessage != False:
                if len(self.submessages[str(p["packetsID"])]["data"]) > (p["packetCount"]-1):
                    self.finishedSubmessages[str(p["packetsID"])] = Message.reassemble(
                        None, self.submessages[str(p["packetsID"])], self.yctx["cryptographyInfo"]["val"], packetCount=p["packetCount"], subMessage=True, yctx=self.yctx
                    )
                    # TODO: Trigger doact
                    return await self.doAct(subMessage=self.finishedSubmessages[str(p["packetsID"])])
            elif len(self.data) >= (self.yctx["packetCount"]["val"]):
                logger.log(
                    20, "Finished receiving for primary message " +
                        str(self.yctx["packetsID"]["val"])
                )
                """
                if self.yctx["wantFullResponse"]["val"] != False:
                    # TO DO: implement loop
                    # responseLoop(packets_id)
                    pass
                """
                # self.data = Message.reassemble(None, self, self.yctx["cryptographyInfo"]["val"])
                return await self.doAct(repeatDataOnActions=rdoaoc)

    async def dump(self):
        msg = {}
        msg["packetsID"] = self.yctx["packetsID"]["val"]
        msg["data"] = self.data
        msg["pAction"] = self.yctx["pAction"]["val"]
        msg["submessages"] = self.submessages
        return msg

    async def doAct(
        self,
        setpAction=False,
        repeatDataOnActions=[],
        subMessage=False
    ):
        # TODO; Get submessage instead of primary
        m = await self.dump()
        if subMessage != False:
            logger.log(30, "Submessage complete")
            # self.submessages[m["packetsID"]] = self.yctx["crytopgraphyInfo"]["val"].decrypt(m["data"])
            # TODO: Debug
            logger.log(10, subMessage)
            if "nonce" in subMessage.keys():
                self.nonce = subMessage["nonce"]
                self.tag = subMessage["tag"]
        if self.checkComplete():
            if len(self.data) >= (self.yctx["packetCount"]["val"]):
                # TODO: Raising out of bounds error
                pAction = self.pActions[int(m["pAction"])]
                logger.log(30, "Full message completed")
                # TODO: Decrypt
                self.data = Message.reassemble(
                    None, self, self.yctx["cryptographyInfo"]["val"])
                """
                data = self.cryptographyInfo.decrypt(m["data"])
                if data == False:
                    self.cryptographyInfo.sessionSetup()
                    data = self.cryptographyInfo.decrypt(m["data"])
                """
                # TODO: Go over whether we need to limit
                # this from being a field in the data for
                # primary messages
                """
                todo.append(
                    Action(
                        "cLog", {"message": "Unknown pAction " + m["pAction"], "priority": 20}
                    )
                )
        )
                """
                self.data["submessages"] = self.submessages
                self.data["yctx"] = self.yctx
                act = Action(pAction, self.data)
                self.todo.append(act)
                if len(repeatDataOnActions) > 0:
                    for caction in repeatDataOnActions:
                        cact = Action(caction, self.data)
                        #self.todo.add(cact)
                        self.todo.append(cact)