piermesh/src/Packets/Message.py

124 lines
3.6 KiB
Python
Executable File

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
class Message:
"""
Full message which is composed of `Packets.Packet.Packet`s
Attributes
----------
packets: list[Packets.Packet.Packet]
List of packets making up the Message
"""
def __init__(
self,
bytesObject: bytes,
sender: int,
senderDisplayName: int,
recipient: int,
recipientNode: int,
dataSize: int = 128,
wantFullResponse: bool = False,
packetsClass: int = 0,
):
"""
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:
bytesObject = lzma.compress(bytesObject)
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,
recipient,
recipientNode,
wantFullResponse=wantFullResponse,
packetsClass=packetsClass,
),
)
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):
"""
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(data))
return res