Files
packetserver/packetserver/common/__init__.py

343 lines
12 KiB
Python

from pe.connect import Connection, ConnectionState
from threading import Lock
from msgpack import Unpacker
from msgpack import packb, unpackb
from enum import Enum
import bz2
from typing import Union, Self
import datetime
import logging
import ax25
class PacketServerConnection(Connection):
connection_subscribers = []
receive_subscribers = []
max_send_size = 2000
def __init__(self, port, call_from, call_to, incoming=False):
super().__init__(port, call_from, call_to, incoming=incoming)
# Now perform any initialization of your own that you might need
self.data = Unpacker()
self.data_lock = Lock()
self.connection_created = datetime.datetime.now(datetime.UTC)
self.connection_last_activity = datetime.datetime.now(datetime.UTC)
self.closing = False
@property
def local_callsign(self):
if self.incoming:
return self.call_to
else:
return self.call_from
@property
def remote_callsign(self):
if self.incoming:
return self.call_from
else:
return self.call_to
def connected(self):
logging.debug("connected")
logging.debug(f"new connection from {self.call_from} to {self.call_to}")
for fn in PacketServerConnection.connection_subscribers:
fn(self)
def disconnected(self):
logging.debug(f"connection disconnected: {self.call_from} -> {self.call_to}")
def data_received(self, pid, data):
self.connection_last_activity = datetime.datetime.now(datetime.UTC)
logging.debug(f"received data: {data}")
with self.data_lock:
logging.debug(f"fed received data to unpacker {data}")
self.data.feed(data)
for fn in PacketServerConnection.receive_subscribers:
logging.debug("found function to notify about received data")
fn(self)
logging.debug("notified function about received data")
def send_data(self, data: Union[bytes, bytearray]):
logging.debug(f"sending data: {data}")
self.connection_last_activity = datetime.datetime.now(datetime.UTC)
if len(data) > self.max_send_size:
logging.debug(f"Large frame detected {len(data)} breaking it up into chunks")
index = 0
counter = 0
while index <= len(data):
logging.debug(f"Sending chunk {counter}")
if (len(data) - index) < self.max_send_size:
super().send_data(data[index:])
break
super().send_data(data[index:index + self.max_send_size])
index = index + self.max_send_size
counter = counter + 1
else:
super().send_data(data)
@classmethod
def query_accept(cls, port, call_from, call_to):
return True
class Message:
"""Base class for communication encapsulated in msgpack objects."""
class CompressionType(Enum):
NONE = 0
BZIP2 = 1
GZIP = 2
DEFLATE = 3
class MessageType(Enum):
REQUEST = 0
RESPONSE = 1
def __init__(self, msg_type: MessageType, compression: CompressionType, payload: dict):
self.type = Message.MessageType(msg_type)
self.compression = Message.CompressionType(compression)
self.data = payload
@property
def vars(self) -> dict:
if 'v' in self.data:
if type(self.data['v']) is dict:
return self.data['v']
return {}
def get_var(self, key: str):
if 'v' not in self.data:
raise KeyError(f"Variable '{key}' not found.")
if str(key).lower() not in self.data['v']:
raise KeyError(f"Variable '{key}' not found.")
return self.data['v'][str(key).lower()]
def set_var(self, key: str, value):
if 'v' not in self.data:
self.data['v'] = {}
self.data['v'][str(key).lower()] = value
@property
def data_bytes(self):
return packb(self.data)
def pack(self) -> bytes:
output = {'t': self.type.value, 'c': self.compression.value}
data_bytes = self.data_bytes
logging.debug("Packing Message")
if (self.compression is self.CompressionType.NONE) or (len(data_bytes) < 30):
output['d'] = data_bytes
output['c'] = self.CompressionType.NONE.value
return packb(output)
if self.compression is self.CompressionType.BZIP2:
compressed = bz2.compress(packb(self.data))
else:
raise NotImplementedError(f"Compression type {self.compression.name} is not implemented yet.")
if len(compressed) < len(data_bytes):
output['d'] = compressed
else:
output['d'] = data_bytes
output['c'] = self.CompressionType.NONE.value
return packb(output)
@property
def payload(self):
if 'd' in self.data:
pl = self.data['d']
if type(pl) in (dict, str, bytes, list):
return pl
else:
return str(pl)
else:
return ""
@payload.setter
def payload(self, payload: Union[str, bytes, dict, list]):
logging.debug(f"Setting a message payload: {type(payload)}: {payload}")
if type(payload) in (str, bytes, dict, list):
logging.debug(f"Payload type is {type(payload)}, conversion to string unnecessary")
self.data['d'] = payload
else:
logging.debug("payload type is not in (str, bytes, dict, list); converting to string")
self.data['d'] = str(payload)
logging.debug(f"Final payload is: {type(payload)}: {payload}")
@classmethod
def partial_unpack(cls, msg: dict) -> Self:
unpacked = msg
comp = Message.CompressionType(unpacked['c'])
msg_type = Message.MessageType(unpacked['t'])
raw_data = unpacked['d']
if comp is Message.CompressionType.NONE:
data = unpackb(raw_data)
elif comp is Message.CompressionType.BZIP2:
data = unpackb(bz2.decompress(raw_data))
else:
raise NotImplementedError(f"Compression type {comp.name} is not implemented yet.")
return Message(msg_type, comp, data)
@classmethod
def unpack(cls, msg_bytes: bytes) -> Self:
try:
unpacked = unpackb(msg_bytes)
except Exception as e:
raise ValueError("ERROR: msg_bytes didn't contain a valid msgpack object.\n" + str(e))
if type(unpacked) is not dict:
raise ValueError("ERROR: unpacked message was not a packetserver message.")
for i in ('t', 'c', 'd'):
if i not in unpacked:
raise ValueError("ERROR: unpacked message was not a packetserver message.")
return Message.partial_unpack(unpacked)
class Request(Message):
class Method(Enum):
GET = 0
POST = 1
UPDATE = 2
DELETE = 3
def __init__(self, msg: Message):
if msg.type is not Message.MessageType.REQUEST:
raise ValueError(f"Can't create a Request Object from a {msg.type} Message object.")
super().__init__(msg.type, msg.compression, msg.data)
if ('p' in msg.data) and (type(msg.data['p']) is not str):
raise ValueError("Path of Request must be a string.")
if 'p' in self.data:
self.data['p'] = str(self.data['p']).strip().lower()
if 'm' in msg.data:
if type(msg.data['m']) is not bytes:
raise ValueError("Method of Request must be bytes.")
self.Method(int(self.data['m'][0]))
@property
def path(self):
if 'p' in self.data:
return str(self.data['p']).lower().strip()
else:
return ""
@path.setter
def path(self, path: str):
self.data['p'] = str(path).strip().lower()
@property
def method(self) -> Method:
if 'm' in self.data:
return self.Method(int(self.data['m'][0]))
else:
return self.Method.GET
@method.setter
def method(self, meth: Method):
meth = self.Method(meth)
self.data['m'] = meth.value.to_bytes(1)
@classmethod
def unpack(cls, msg_bytes: bytes) -> Self:
msg = super().unpack(msg_bytes)
return Request(msg)
@classmethod
def blank(cls) -> Self:
msg = Message(Message.MessageType.REQUEST, Message.CompressionType.NONE, {})
return Request(msg)
def __repr__(self):
return f"<Request: {self.method.name} '{self.path}'>"
class Response(Message):
def __init__(self, msg: Message):
if msg.type is not Message.MessageType.RESPONSE:
raise ValueError(f"Can't create a Response Object from a {msg.type} Message object.")
super().__init__(msg.type, msg.compression, msg.data)
if 'c' in msg.data:
status_bytes = self.data['c']
if type(status_bytes) is not bytes:
raise ValueError("Invalid Response data")
status_code = int.from_bytes(status_bytes)
if status_code >= 600:
raise ValueError("Invalid status code.")
@classmethod
def unpack(cls, msg_bytes: bytes) -> Self:
msg = super().unpack(msg_bytes)
return Response(msg)
@classmethod
def blank(cls) -> Self:
msg = Message(Message.MessageType.RESPONSE, Message.CompressionType.NONE, {})
return Response(msg)
@property
def status_code(self) -> int:
if 'c' in self.data:
status_bytes = self.data['c']
if type(status_bytes) is not bytes:
raise ValueError("Invalid Response data")
status_code = int.from_bytes(status_bytes)
if status_code >= 600:
raise ValueError("Invalid status code.")
return status_code
else:
return 200
@status_code.setter
def status_code(self, code: int):
if (code <= 0) or (code >= 600):
raise ValueError("Status must be a positive integer <= 600")
self.data['c'] = code.to_bytes(2)
def __repr__(self):
return f"<Response: {self.status_code}>"
def send_response(conn: PacketServerConnection, response: Response, original_request: Request,
compression: Message.CompressionType = Message.CompressionType.BZIP2):
if conn.state.name == "CONNECTED" and not conn.closing:
# figure out compression setting based on request
logging.debug("Determining compression of response")
comp = compression
logging.debug(f"Default comp: {comp}")
logging.debug(f"Original vars: {original_request.vars}")
if 'C' in original_request.vars:
logging.debug(f"Detected compression header in original request: {original_request.vars['C']}")
val = original_request.vars['C']
for i in Message.CompressionType:
logging.debug(f"Checking type: {i}")
if str(val).strip().upper() == i.name:
comp = i
logging.debug(f"matched compression with var to {comp}")
break
try:
if int(val) == i.value:
comp = i
logging.debug(f"matched compression with var to {comp}")
except ValueError:
pass
response.compression = comp
logging.debug(f"Final compression: {response.compression}")
logging.debug(f"sending response: {response}, {response.compression}, {response.payload}")
conn.send_data(response.pack())
logging.debug("response sent successfully")
else:
logging.warning(f"Attempted to send data, but connection state is {conn.state.name}")
def send_blank_response(conn: PacketServerConnection, original_request: Request, status_code: int = 200,
payload: Union[bytes, bytearray, str, dict, list] = ""):
response = Response.blank()
response.status_code = status_code
response.payload = payload
send_response(conn, response, original_request)