sampy_archive/sampy/client.py

import socket
import logging
import struct

from .network.compression import uncompress
from .handlers import serverQueryHandler, mainHandler
from .helpers.byteFormater import readable_bytes
from .server import Server

from . import glob
if glob.config["debug"]:
	from importlib import reload

logger = logging.getLogger(__name__)

class Client:
	BUFF_SIZE = 65535
	HANDLERS = {
		b"SAMP": serverQueryHandler.handle
	}

	def __init__(self, server, addr: "Address"):
		self.server = server
		
		self.ip_uint = struct.unpack(b"<I", bytes(int(x) for x in addr[0].split(".")))[0]

		self.addr = addr
		self.state = 0

		self.__writebuffer = bytearray()
	
	def handle_data(self, data: bytes):
		found = False
		for pattern, handle in self.HANDLERS.items(): # SAMP query
			if data.startswith(pattern):
				queue = handle(self, data)
				if queue:
					self.write(queue)
				found = True

		if not found:
			# deobfuscate
			try:
				data: bytearray = uncompress(data)
				logger.debug("[%s] -> %s" % (self.addr, readable_bytes(data)))
				queue = mainHandler.handle(self, data)
			except Exception as err:
				logger.error(err)
			finally:
				if queue:
					self.write(queue)
	
	def write(self, data: bytes):
		if type(data) == int:
			self.__writebuffer.append(data)
		elif type(data) == list:
			for item in data:
				self.write(item)
		else:
			self.__writebuffer += data

	def get_buffer_size(self) -> int:
		return len(self.__writebuffer)
	
	def flush(self):
		if not self.get_buffer_size():
			return

		logger.debug("[%s] <- %s" % (self.addr, readable_bytes(self.__writebuffer)))
		self.server.socket.sendto(self.__writebuffer, self.addr)
		self.__writebuffer.clear()

	def send_immediate(self, data: bytes):
		logger.debug("[%s] <- %s" % (self.addr, readable_bytes(data)))
		self.server.socket.sendto(data, self.addr)

def restart():
	reload(serverQueryHandler)
	reload(mainHandler)
	mainHandler.restart()