Massive refactoring addition of API

5 files changed, 112 insertions, 37 deletions

diff --git a/main.py b/main.py
index 25ae922..dc80db6 100644
--- a/main.py
+++ b/main.py
@@ -5,12 +5,13 @@ The mighty silly webserver written in python for no good reason
 
 import ssl
 import gzip
-import time
 import socket
 import brotli
 import signal
 import asyncio
 import aiofiles
+from src import APIv1
+from src.socks import *
 from src.request import Request
 from src.minimizer import minimize_html
 
@@ -40,6 +41,14 @@ PATH_MAP = {
          "compress": True},
 }
 
+# API
+API_PATH = {
+    "/APIv1/file/generated/1gib",
+    "/APIv1/file/generated/5gib",
+    "/APIv1/file/generated/10gib",
+    "/APIv1/file/generated/20gib",
+}
+
 # internal path map
 I_PATH_MAP = {
     "/err/response.html":       {"path": "www/err/response.html"}
@@ -122,10 +131,10 @@ class HTTPServer:
         while True:
             # try to accept new connection
             try:
-                client = (await self._accept(self.socket))[0]
+                client = (await ssl_sock_accept(self.socket))[0]
 
             # if socket was closed -> break
-            except OSError:
+            except OSError as e:
                 print("Closed.")
                 break
 
@@ -172,8 +181,6 @@ class HTTPServer:
     async def handle_get_request(client: ssl.SSLSocket, request: Request):
         """
         Handles user's GET request
-        :param client: client
-        :param request: client's request
         """
 
         # get available compression methods
@@ -200,19 +207,47 @@ class HTTPServer:
 
             # send 200 response with the file to the client
             await HTTPServer._send(client, 200, data, headers)
-        else:
-            # in case of error, return error page
-            async with aiofiles.open(I_PATH_MAP["/err/response.html"]["path"], "r") as f:
-                data = await f.read()
 
-            # status code
-            status_code = 404
+            # return after answer
+            return
+
+        # if it's an API request
+        elif request.path in API_PATH:
+            # get API version
+            api_version = request.path.split("/")[1]
+
+            match api_version:
+                case "APIv1":
+                    status, data = await APIv1.respond(client, request)
+                case _:
+                    status = 400
 
-            # format error response
-            data = data.format(status_code=get_response_code(status_code).decode("ascii"))
+            # if status is not 200 -> send bad response
+            if status != 200:
+                await HTTPServer._bad_response(client, status)
 
-            # send 404 response to the client
-            await HTTPServer._send(client, status_code, data.encode("ascii"))
+            # return after answer
+            return
+
+        # in case of error, return error page
+        await HTTPServer._bad_response(client, 404)
+
+    @staticmethod
+    async def _bad_response(client: ssl.SSLSocket, status_code: int):
+        """
+        Sends a bad response page to the client.
+        :param client: client
+        :param status_code: status code
+        """
+
+        async with aiofiles.open(I_PATH_MAP["/err/response.html"]["path"], "r") as f:
+            data = await f.read()
+
+        # format error response
+        data = data.format(status_code=get_response_code(status_code).decode("ascii"))
+
+        # send response to the client
+        await HTTPServer._send(client, status_code, data.encode("ascii"))
 
     @staticmethod
     async def _send(client: ssl.SSLSocket, response: int, data: bytes = None, headers: dict[str, str] = None):
@@ -245,7 +280,7 @@ class HTTPServer:
             byte_header += f"{key}: {value}\r\n".encode("ascii")
 
         # send response to the client
-        await HTTPServer._sendall(
+        await ssl_sock_sendall(
             client,
             b'HTTP/1.1 ' +
             get_response_code(response) +
@@ -276,7 +311,7 @@ class HTTPServer:
         while True:
             try:
                 # fetch packet
-                message = await self._recv(client, self.packet_size)
+                message = await ssl_sock_recv(client, self.packet_size)
             except OSError:
                 break
 
@@ -295,24 +330,6 @@ class HTTPServer:
         # return empty buffer on error
         return b''
 
-    @staticmethod
-    async def _accept(sock: ssl.SSLSocket) -> tuple[ssl.SSLSocket, str]:
-        while True:
-            try:
-                return sock.accept()
-            except BlockingIOError:
-                time.sleep(1.e-3)
-
-    @staticmethod
-    async def _recv(sock: ssl.SSLSocket, buflen: int = 1024):
-        while (msg := sock.recv(buflen)) == b'':
-            time.sleep(1.e-3)
-        return msg
-
-    @staticmethod
-    async def _sendall(sock: ssl.SSLSocket, data: bytes):
-        sock.sendall(data)
-
 
 def main():
     server = HTTPServer(port=13700)
diff --git a/requirements.txt b/requirements.txt
index 2d404ff..07bfe5c 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,2 +1,2 @@
 aiofiles==24.1.0
-htmlmin==0.1.12
+Brotli~=1.1.0
\ No newline at end of file
diff --git a/src/APIv1.py b/src/APIv1.py
new file mode 100644
index 0000000..e0e4848
--- /dev/null
+++ b/src/APIv1.py
@@ -0,0 +1,34 @@
+import random
+from src.socks import *
+from ssl import SSLSocket
+from src.request import Request
+
+
+async def respond(client: SSLSocket, request: Request) -> tuple[int, bytes]:
+    """
+    Respond to clients API request
+    """
+
+    # decode API request
+    split_path = request.path.split("/")
+    api_level1 = split_path[2]
+    api_level2 = split_path[3]
+    api_request = split_path[4]
+
+    # do something with it (oh god)
+    if api_level1 == "file":
+        if api_level2 == "generated":
+            if api_request == "1gib":
+                return 200, random.randbytes(2**30 * 1)
+            elif api_request == "5gib":
+                return 200, random.randbytes(2**30 * 5)
+            elif api_request == "10gib":
+                return 200, random.randbytes(2**30 * 10)
+            elif api_request == "20gib":
+                return 200, random.randbytes(2**30 * 20)
+            else:
+                return 400, b''
+        else:
+            return 400, b''
+    else:
+        return 400, b''
diff --git a/src/socks.py b/src/socks.py
new file mode 100644
index 0000000..119a62a
--- /dev/null
+++ b/src/socks.py
@@ -0,0 +1,20 @@
+import time
+from ssl import SSLSocket
+
+
+async def ssl_sock_accept(sock: SSLSocket) -> tuple[SSLSocket, str]:
+    while True:
+        try:
+            return sock.accept()
+        except BlockingIOError:
+            time.sleep(1.e-3)
+
+
+async def ssl_sock_recv(sock: SSLSocket, buflen: int = 1024):
+    while (msg := sock.recv(buflen)) == b'':
+        time.sleep(1.e-3)
+    return msg
+
+
+async def ssl_sock_sendall(sock: SSLSocket, data: bytes):
+    sock.sendall(data)
diff --git a/www/test.html b/www/test.html
index 9e37f1e..e3befbf 100644
--- a/www/test.html
+++ b/www/test.html
@@ -11,7 +11,11 @@
     </header>
     <div id="section-div">
         <section>
-
+            <h1> File download test </h1>
+            <p> <a href="/APIv1/file/generated/1gib" target="_blank">Download 1 GiB file</a> </p>
+            <p> <a href="/APIv1/file/generated/5gib" target="_blank">Download 5 GiB file</a> </p>
+            <p> <a href="/APIv1/file/generated/10gib" target="_blank">Download 10 GiB file</a> </p>
+            <p> <a href="/APIv1/file/generated/20gib" target="_blank">Download 20 GiB file</a> </p>
         </section>
     </div>
     <footer>