sampy_archive/objects/client.py

from importlib import reload

import socket
import logging
import struct

from network.compression import uncompress

from handlers import serverQueryHandler, mainHandler

from helpers.byteFormater import readable_bytes

logger = logging.getLogger(__name__)

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

	def __init__(self, server: "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_to_buffer(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_to_buffer(queue)
	
	def write_to_buffer(self, data: bytes):
		if type(data) == int:
			self.__writebuffer.append(data)
		elif type(data) == list:
			for item in data:
				self.write_to_buffer(item)
		else:
			self.__writebuffer += data

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

	def write_socket(self):
		if self.get_buffer_size() > 0:
			#write_chunk = self.__writebuffer.pop(0)
			#logger.debug("[%s] <- %s" % (self.addr, write_chunk))
			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()