Merge pull request #1 from UltraQbik/rewrite

Rewrite

1 files changed, 192 insertions, 187 deletions

diff --git a/main.py b/main.py
index 4078d99..4f82213 100644
--- a/main.py
+++ b/main.py
@@ -8,255 +8,260 @@ import json
 import gzip
 import socket
 import asyncio
-import htmlmin
 import aiofiles
 import threading
 
 
-# some constants
-PACKET_SIZE = 2048
-PORT = 13700            # using random port cuz why not
-
-
-# response status codes
-RESPONSE = {
-    200: b'OK',
-    400: b'Bad Request',
-    401: b'Unauthorized',
-    403: b'Forbidden',
-    404: b'Not Found',
-    6969: b'UwU'
+# path mapping
+PATH_MAP = {
+    "/":                        {"path": "www/index.html"},
+    "/index.html":              {"path": "www/index.html"},
+    "/robots.txt":              {"path": "www/robots.txt"},
+    "/favicon.ico":             {"path": "www/favicon.ico"},
+    "/css/styles.css":          {"path": "css/styles.css"},
 }
 
 
-def get_response(code: int) -> bytes:
-    return str(code).encode("ascii") + RESPONSE.get(code, b':(')
-
-
-def is_alive(sock: socket.socket) -> bool:
+def get_response_code(code: int) -> bytes:
+    match code:
+        case 200:
+            return b'200 OK'
+        case 400:
+            return b'400 Bad Request'
+        case 401:
+            return b'401 Unauthorized'
+        case 403:
+            return b'403 Forbidden'
+        case 404:
+            return b'404 Not Found'
+        case 6969:
+            return b'6969 UwU'
+        case _:  # in any other case return bad request response
+            return get_response_code(400)
+
+
+class Request:
     """
-    Checks if the socket is still alive
-    :param sock: socket
-    :return: boolean (true if socket is alive, false otherwise)
+    Just a request
     """
-    return getattr(sock, '_closed', False)
 
+    def __init__(self):
+        self.type: str = ""
+        self.path: str = ""
 
-def decode_request(req: str) -> dict[str, str | list | None]:
-    # request dictionary
-    request = dict()
+    @staticmethod
+    def create(raw_request: bytes):
+        """
+        Creates self class from raw request
+        :param raw_request: bytes
+        :return: self
+        """
 
-    # request type and path
-    request["type"] = req[:6].split(" ")[0]
-    request["path"] = req[len(request["type"]) + 1:req.find("\r\n")].split(" ")[0]
+        # new request
+        request = Request()
 
-    # decode other headers
-    for line in req.split("\r\n")[1:]:
-        if len(split := line.split(":")) == 2:
-            key, value = split
-            value = value.lstrip(" ")
+        # fix type and path
+        request.type = raw_request[:raw_request.find(b' ')].decode("ascii")
+        request.path = raw_request[len(request.type)+1:raw_request.find(b' ', len(request.type)+1)].decode("ascii")
 
-            # write key value pair
-            request[key] = value
+        # decode headers
+        for raw_header in raw_request.split(b'\r\n'):
+            if len(pair := raw_header.decode("ascii").split(":")) == 2:
+                key, val = pair
+                val = val.strip()
 
-    return request
+                # set attribute to key value pair
+                setattr(request, key, val)
 
+        # return request
+        return request
 
-class HTMLServer:
+    def __str__(self):
+        return '\n'.join([f"{key}: {val}" for key, val in self.__dict__.items()])
+
+
+class HTTPServer:
     """
-    The very cool webserver
+    The mighty HTTP server
     """
 
-    def __init__(self):
-        self.sock: socket.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        self.clients: list[socket.socket] = []
+    def __init__(self, *, port: int, packet_size: int = 2048):
+        self.bind_port: int = port
+        self.packet_size: int = packet_size
+        self.socket: socket.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 
-        # list of allowed paths
-        self.allowed_paths: dict[str, dict] = {
-            "/":                {"path": "www/index.html",      "encoding": "css/html"},
-            "/robots.txt":      {"path": "www/robots.txt",      "encoding": "text"},
-            "/favicon.ico":     {"path": "www/favicon.ico",     "encoding": "bin"},
-            "/css/styles.css":  {"path": None,                  "encoding": "css/html"},
-        }
+        self.clients: list[socket.socket] = []
 
-    def run(self):
+    def start(self):
         """
-        Function that starts the webserver
+        Method to start the web server
         """
 
-        # bind the server to port and start listening
-        self.sock.bind(('', PORT))
-        self.sock.listen()
+        # bind and start listening to port
+        self.socket.bind(('', self.bind_port))
+        self.socket.listen()
 
-        # start running thread
-        t = threading.Thread(target=self._run, daemon=True)
-        t.start()
+        # start the listening thread
+        threading.Thread(target=self._listen_thread, daemon=True).start()
 
         # keep alive
         try:
             while True:
+                # sleep 100 ms, otherwise the while true will 100% one of your cores
                 time.sleep(0.1)
+
+        # shutdown on keyboard interrupt
         except KeyboardInterrupt:
-            self.sock.close()
+            self.socket.close()
             print("Closed.")
 
-    def _run(self):
+    def _listen_thread(self):
         """
-        Run function for threads
-        :return:
+        Listening for new connections
         """
 
-        asyncio.run(self.server_listener())
-
-    async def server_listener(self):
-        """
-        Listens for new connections, and handles them
-        """
+        # run the coroutine
+        asyncio.run(self._thread_listen_coro())
 
+    async def _thread_listen_coro(self):
         while True:
-            client, address = self.sock.accept()
+            # accept new connection, add to client list and start listening to it
+            client, _ = self.socket.accept()
             self.clients.append(client)
-            await self.server_handle(client)
+            await self.client_handle(client)
 
-    async def server_handle(self, client: socket.socket):
+    async def client_handle(self, client: socket.socket):
         """
-        Handles the actual connections (clients)
-        :param client: connection socket
+        Handles client's connection
         """
 
-        # message buffer
-        buffer = bytearray()
         while True:
-            # try to fetch a message
-            # die otherwise
+            # receive request from client
+            raw_request = self._recvall(client)
+
+            # decode request
+            request: Request = Request.create(raw_request)
+
+            # log request
+            async with aiofiles.open("logs.log", "a") as f:
+                await f.write(f"IP: {client.getpeername()[0]}\n{request}\n\n")
+
+            # handle requests
             try:
-                message = client.recv(PACKET_SIZE)
-            except OSError:
-                break
-            if message == b'':
+                match request.type:
+                    case "GET":
+                        await self.handle_get_request(client, request)
+                    case _:
+                        break
+
+            # break on exception
+            except Exception as e:
+                print(e)
                 break
 
-            # append packet to buffer
-            buffer += message
+            # break the connection
+            break
 
-            # check EoF (2 blank lines)
-            if buffer[-4:] == b'\r\n\r\n':
-                # text buffer
-                text_buffer = buffer.decode("ascii")
-
-                # decode request
-                request = decode_request(text_buffer)
-
-                print(f"[{request['type']}] Request from client '{client.getpeername()[0]}'")
-
-                # log that request
-                async with aiofiles.open("logs.log", "a") as f:
-                    await f.write(
-                        json.dumps(
-                            {
-                                "client": client.getpeername()[0],
-                                "request": request
-                            },
-                            indent=2
-                        ) + "\n"
-                    )
-
-                # handle the request
-                if request["type"] == "GET":
-                    await self.handle_get_request(client, request)
-                else:
-                    await self.handle_other_request(client)
-
-                # clear buffer
-                buffer.clear()
-        client.close()
-        self.clients.remove(client)
-
-    async def handle_get_request(self, client: socket.socket, req: dict[str, str | None]):
-        # check if the path is too long
-        if len(req["path"]) > 255:
-            response = get_response(400)
-            data = b''
-
-        # if it's yandex
-        elif req.get("from") == "support@search.yandex.ru":
-            response = get_response(404)
-            data = b'Nothing...'
-
-        # check UwU path
-        elif req["path"] == "/UwU":
-            response = get_response(6969)
-            data = b'<h1>' + b'UwU ' * 2000 + b'</h1>'
-
-        # otherwise check access
-        elif req["path"] in self.allowed_paths:
-            # get path
-            path = self.allowed_paths[req["path"]]["path"]
-
-            # if path is None, return generic filepath
-            if path is None:
-                path = req["path"][1:]
-
-            # check encoding
-            if self.allowed_paths[req["path"]]["encoding"] == "css/html":
-                # return text data
-                async with aiofiles.open(path, "r") as f:
-                    data = htmlmin.minify(await f.read()).encode("ascii")
-            else:
-                # return binary / text data
-                async with aiofiles.open(path, "rb") as f:
-                    data = await f.read()
-            response = get_response(200)
-
-        # in any other case
-        else:
-            response = get_response(403)
-            data = b'Idk what you are trying to do here :/'
+        # close connection (stop page loading)
+        self._close_client(client)
 
-        # make headers
-        headers = {}
+    @staticmethod
+    async def handle_get_request(client: socket.socket, request: Request):
+        """
+        Handles user's GET request
+        :param client: client
+        :param request: client's request
+        """
 
-        # check if compression is supported
-        if req.get("Accept-Encoding"):
-            encoding_list = [enc.lstrip(" ") for enc in req["Accept-Encoding"].split(",")]
+        # check if request path is in the PATH_MAP
+        if request.path in PATH_MAP:
+            # if it is -> return file from that path
+            async with aiofiles.open(PATH_MAP[request.path]["path"], "rb") as f:
+                data = await f.read()
 
-            # check for gzip, and add to headers if present
-            if "gzip" in encoding_list:
-                headers["Content-Encoding"] = "gzip"
+            # send 200 response with the file to the client
+            HTTPServer._send(client, 200, data)
+        else:
+            # send 400 response to the client
+            HTTPServer._send(client, 400)
 
-        # send response
-        await self.send(client, response, data, headers)
-        client.close()
+    @staticmethod
+    def _send(client: socket.socket, response: int, data: bytes = None, headers: dict[str, str] = None):
+        """
+        Sends client response code + headers + data
+        :param client: client
+        :param response: response code
+        :param data: data
+        :param headers: headers to include
+        """
+
+        # if data was not given
+        if data is None:
+            data = bytes()
+
+        # if headers were not given
+        if headers is None:
+            headers = dict()
+
+        # format headers
+        byte_header = bytearray()
+        for key, value in headers.items():
+            byte_header += f"{key}: {value}\r\n".encode("ascii")
+
+        # send response to the client
+        client.sendall(
+            b'HTTP/1.1 ' +
+            get_response_code(response) +
+            b'\r\n' +
+            byte_header +       # if empty, we'll just get b'\r\n\r\n'
+            b'\r\n' +
+            data
+        )
+
+    def _close_client(self, client: socket.socket):
+        """
+        Closes a client
+        """
 
-    async def handle_other_request(self, client: socket.socket):
-        # just say 'no'
-        await self.send(client, get_response(403), b'No. Don\'t do that, that\'s cringe')
         client.close()
+        if client in self.clients:
+            self.clients.remove(client)
+
+    def _recvall(self, client: socket.socket) -> bytes:
+        """
+        Receive All (just receives the whole message, instead of 1 packet at a time)
+        """
 
-    async def send(self, client: socket.socket, response: bytes, data: bytes, headers: dict[str, str] | None = None):
-        # construct headers
-        formatted_headers = b''
-        if headers is not None:
-            formatted_headers = "".join([f"{key}: {val}\r\n" for key, val in headers.items()]).encode("ascii")
+        # create message buffer
+        buffer: bytearray = bytearray()
 
-            # check for compression
-            if headers.get("Content-Encoding") == "gzip":
-                # compress data
-                data = gzip.compress(data)
+        # start fetching the message
+        while True:
+            try:
+                # fetch packet
+                message = client.recv(self.packet_size)
+            except OSError:
+                break
 
-        # construct message
-        if formatted_headers == b'':
-            message = b'HTTP/1.1 ' + response + b'\r\n\r\n' + data
-        else:
-            message = b'HTTP/1.1 ' + response + b'\r\n' + formatted_headers + b'\r\n' + data
+            # that happens when user stops loading the page
+            if message == b'':
+                break
+
+            # append fetched message to the buffer
+            buffer += message
+
+            # check for EoF
+            if buffer[-4:] == b'\r\n\r\n':
+                # return the received message
+                return buffer
 
-        # send message to client
-        client.sendall(message)
+        # return empty buffer on error
+        return b''
 
 
 def main():
-    server = HTMLServer()
-    server.run()
+    server = HTTPServer(port=13700)
+    server.start()
 
 
 if __name__ == '__main__':