server.py

import sys
import socket
from database import Database
from command import handle_command
from request import Request, RequestType
from response import Response, SystemResponse, ExtensionResponse
from objects.cafe import Cafe
from typing import TYPE_CHECKING
import asyncio
import multiprocessing

if TYPE_CHECKING:
    from objects.player import Player

SERVER_HOST = '127.0.0.1'
SERVER_PORT = 9339
MAX_CONN = 10


class CafeServer:
    def __init__(self) -> None:
        self.host: str = SERVER_HOST
        self.port: int = SERVER_PORT
        self.max_conn: int = MAX_CONN
        self.running: bool = False
        self.clients: dict[tuple, 'Player'] = {}
        self.players: list['Player'] = []
        self.socket: socket.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

        self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.bind_socket()

        self.loop = asyncio.new_event_loop()
        self.process = multiprocessing.Process(target=self.run_server)

        self.db = Database()
        self.marketplace = Cafe.get_marketplace()

    @staticmethod
    async def send_response(connection, response: Response):
        response = response.raw_response
        connection.write(response.encode())
        await connection.drain()

    def bind_socket(self) -> None:
        try:
            self.socket.bind((self.host, self.port))
        except PermissionError:
            sys.stderr.write(f'[ERROR] You do not have permission to use {self.host}:{self.port}.\n')
            sys.exit(1)
        except OSError:
            sys.stderr.write(f'[ERROR] {self.host}:{self.port} is already in use.\n')
            sys.exit(1)

    async def handle_client(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
        tasks = set()
        peername = writer.get_extra_info('peername')

        print(f'[CONNECTED] {peername}')

        try:
            while not writer.is_closing():
                try:
                    request_str = (await reader.read(1024)).decode()
                except ConnectionResetError:
                    break

                if not request_str:
                    break

                try:
                    request = Request(request_str.strip('\x00'))
                    print(f"[RECEIVED] {peername}: {request.raw_request}")
                except Exception as e:
                    print(f"[ERROR] {peername}: {e}")
                    break

                match request.type:
                    case RequestType.DISCONNECT:
                        writer.close()
                        await writer.wait_closed()

                    case RequestType.POLICY_FILE_REQUEST:
                        response = SystemResponse(request.cmd)
                        await self.send_response(writer, response)
                        writer.close()
                        await writer.wait_closed()

                    case RequestType.XML if request.cmd != 'login':
                        response = SystemResponse(request.cmd)
                        await self.send_response(writer, response)

                    case _ if request.cmd != 'pin':
                        task = asyncio.create_task(handle_command(self, writer, request.cmd, *request.params))
                        tasks.add(task)
                        task.add_done_callback(tasks.discard)

        except Exception as e:
            print(f"[ERROR] {peername}: {e}")

        finally:
            player = self.clients.get(peername)
            if player:
                if player.room:
                    if writer in player.room.clients:
                        player.room.clients.remove(writer)

                    for client in player.room.clients:
                        await self.send_response(client, ExtensionResponse('juq', '-1', '0', str(player.id)))

                for waiter in player.cafe.waiters:
                    if waiter.task:
                        print(f'[CANCELLED] {waiter.task}')
                        waiter.task.cancel()

                if player.customer_cycle: # Binded to the Cafe in the NPC OPTIMIZATION BRANCH.
                    player.customer_cycle.cancel() # Binded to the Cafe in the NPC OPTIMIZATION BRANCH.

                player.customer_cycle = None # Binded to the Cafe in the NPC OPTIMIZATION BRANCH.

                for customer in player.cafe.customers:
                    if customer.task:
                        print(f'[CANCELLED] {customer.task}')
                        customer.task.cancel()

                for task in list(tasks):
                    if not task.done():
                        task.cancel()
                        await asyncio.gather(task, return_exceptions=True)

                print(f'[DISCONNECTED] {peername}')
                self.players.remove(player)
                del self.clients[peername]

    def start(self) -> None:
        self.process.start()
        print(f"Server started in process with PID: {self.process.pid}")

        try:
            self.process.join()
        except KeyboardInterrupt:
            print("[INFO] KeyboardInterrupt received. Stopping server...")
            self.stop()

    def run_server(self) -> None:
        asyncio.set_event_loop(self.loop)
        self.running = True

        server = self.loop.run_until_complete(
            asyncio.start_server(
                self.handle_client,
                host=self.host,
                port=self.port,
                reuse_address=True,
                reuse_port=True,
            )
        )

        print(f'[INFO] Listening on port {self.port}...')

        try:
            self.loop.run_until_complete(server.serve_forever())
        except KeyboardInterrupt:
            self.stop()

    def stop(self) -> None:
        self.running = False
        self.socket.close()
        print('[INFO] Server stopped.')