diff --git a/src/core/tcp_server.py b/src/core/tcp_server.py index 6424e1f..e0a06ac 100644 --- a/src/core/tcp_server.py +++ b/src/core/tcp_server.py @@ -18,26 +18,49 @@ class TCPServer: self.host = host self.port = port - async def send(self, data, sync): - pass - async def recv(self, client): - not_alive = client.is_disconnected() - if not not_alive: - self.log.debug(f"Client with ID {client.cid} disconnected") - return "" + if not client.is_disconnected(): + self.log.debug(f"Client with {client.nick}({client.cid}) disconnected") + return + + header = await client.loop.sock_recv(client.socket, 4) # header: 4 bytes data = b"" while True: - chunk = await client.loop.sock_recv(client.socket, 10) + chunk = await client.loop.sock_recv(client.socket, 1) if not chunk: break data += chunk + + int_header = 0 + for i in range(len(header)): + int_header += header[i] + self.log.debug(f"header: `{header}`; int_header: `{int_header}`; data: `{data}`;") + + if int_header <= 0: + client.kick("Invalid packet - header negative") + return + + if not int_header < 100 * MB: + client.kick("Header size limit exceeded") + self.log.warn(f"Client {client.nick}({client.cid}) sent header of >100MB - " + f"assuming malicious intent and disconnecting the client.") + return + + # TODO: read Data + if len(data) != int_header: + self.log.debug(f"WARN Expected to read {int_header} bytes, instead got {len(data)}") + + # TODO: ABG: DeComp(Data) + abg = b"ABG:" + if len(data) > len(abg) and data.startswith(abg): + data = data[len(abg):] + # return DeComp(Data); return data async def auth_client(self, sock): # TODO: Authentication client = self.Core.create_client(sock) - self.log.debug(f"Client: \"IP: {client.addr!r}; ID: {client.cid}\" - started authentication!") + self.log.debug(f"Client: \"IP: {client.addr!r}; ID: {client.cid}\" - Authentication!") data = await self.recv(client) self.log.debug(f"recv1 data: {data}") if len(data) > 50: @@ -47,48 +70,50 @@ class TCPServer: client.kick("Outdated Version.") return else: - self.log.debug('tcp_send(b"A")') - client.tcp_send(b"A") + pass + # self.log.debug('tcp_send(b"A")') + # client.tcp_send(b"A") - data = await self.recv(client) - self.log.debug(f"recv2 data: {data}") + # data = await self.recv(client) + # self.log.debug(f"recv2 data: {data}") client.kick("TODO Authentication") + return False + + async def handle_download(self, sock): + # TODO: HandleDownload + self.log.debug(f"Client: \"IP: {0!r}; ID: {0}\" - HandleDownload!") + return False + + async def handle_code(self, code, sock): + match code: + case "C": + return await self.auth_client(sock) + case "D": + return await self.handle_download(sock) + case "P": + sock.send(b"P") + return True + case _: + self.log.error(f"Unknown code: {code}") + return False async def handle_client(self, sock): - while True: try: data = sock.recv(1) if not data: break - message = data.decode("utf-8").strip() - addr = sock.getsockname() - self.log.debug(f"Received {message!r} from {addr!r}") - code = message[0] - match code: - case "C": - await self.auth_client(sock) - case "D": - # TODO: HandleDownload - print("TODO: HandleDownload") - case "P": - # TODO: Понять что это и зачем... - sock.sendall(b"P") - case _: - self.log.error(f"Unknown code: {code}") + code = data.decode() + self.log.debug(f"Received {code!r} from {sock.getsockname()!r}") + if not await self.handle_code(code, sock): + break except Exception as e: - print("Error:", e) + self.log.error(f"Error: {e}") traceback.print_exc() break - print("Error while connecting..") - - # async def start(self): - # self.log.debug("Starting TCP server.") - # server = await asyncio.start_server(self.handle_client, self.host, self.port, family=socket.AF_INET) - # self.log.debug(f"Serving on {server.sockets[0].getsockname()}") - # async with server: - # await server.serve_forever() + sock.close() + self.log.error("Error while connecting..") async def start(self): self.log.debug("Starting TCP server.") diff --git a/src/core/tcp_server.pyi b/src/core/tcp_server.pyi index 699b69e..1e50f07 100644 --- a/src/core/tcp_server.pyi +++ b/src/core/tcp_server.pyi @@ -17,14 +17,10 @@ class TCPServer: self.Core = core self.host = host self.port = port - - async def send(self, data, sync) -> None: ... - - async def recv(self, writer: Client) -> bytes: ... - + async def recv(self, client: Client) -> bytes: ... async def auth_client(self, sock: socket.socket) -> None: ... - + async def handle_download(self, sock: socket.socket) -> None: ... + async def handle_code(self, code: str, sock: socket.socket) -> None: ... async def handle_client(self, sock: socket.socket) -> None: ... - async def start(self) -> None: ...