"header" + handle_code()

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.")

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: ...