1 files changed, 58 insertions, 0 deletions
diff --git a/src/controller.py b/src/controller.py
new file mode 100644
index 0000000..351d0bc
--- /dev/null
+++ b/src/controller.py
@@ -0,0 +1,58 @@
+import io
+import logging
+import os
+import sanic
+import sqlite3
+import tempfile
+
+import db
+import templates
+import utils
+
+conn = sqlite3.connect('db.sqlite3')
+files_directory = 'files'
+authorized_key = os.environ['KEY']
+
+def index():
+    return sanic.html(templates.index)
+
+async def upload(request):
+    key = request.headers.get('X-Key')
+    if not key == authorized_key:
+        sanic.log.logging.info('Unauthorized to upload file: wrong key')
+        return sanic.text('Unauthorized', status = 401)
+    else:
+        sanic.log.logging.info('Uploading file')
+        content_length = int(request.headers.get('content-length'))
+        filename = utils.sanitize_filename(request.headers.get('X-FileName'))
+        expiration = request.headers.get('X-Expiration')
+
+        with tempfile.NamedTemporaryFile(delete = False) as tmp:
+            while data := await request.stream.read():
+                tmp.write(data)
+
+        sanic.log.logging.info('File uploaded')
+        file_id = db.insert_file(conn, filename, expiration, content_length)
+        os.makedirs(files_directory, exist_ok=True)
+        os.rename(tmp.name, os.path.join(files_directory, file_id))
+
+        return sanic.text(file_id)
+
+async def file(file_id: str, download: bool):
+    res = db.get_file(conn, file_id)
+    if res is None:
+        self._serve_str(templates.not_found, 404, 'text/html')
+    else:
+        filename, expires, content_length = res
+        disk_path = os.path.join(files_directory, file_id)
+        if download:
+            return await sanic.response.file_stream(
+                disk_path,
+                chunk_size = io.DEFAULT_BUFFER_SIZE,
+                headers = {
+                    'Content-Disposition': f'attachment; filename={filename}',
+                    'Content-Length': content_length
+                }
+            )
+        else:
+            return sanic.html(templates.file_page(file_id, filename, expires))