File extension checker (#77)

Author: grigoriev-semyon

calendar_backend/methods/utils.py

@@ -2,9 +2,10 @@
 import os
 import random
 import string
+from typing import Final
 
 import aiofiles
-from fastapi import File, UploadFile
+from fastapi import File, UploadFile, HTTPException
 from sqlalchemy.orm import Session
 
 from calendar_backend.models.db import (
@@ -68,11 +69,14 @@ async def get_lecturer_lessons_in_daterange(
             events_list.append(lesson)
     return events_list
 
+SUPPORTED_FILE_EXTENSIONS: Final[list[str]] = ['png', 'svg', 'jpg', 'jpeg']
 
 async def upload_lecturer_photo(lecturer_id: int, session: Session, file: UploadFile = File(...)) -> Photo:
     lecturer = Lecturer.get(lecturer_id, session=session)
     random_string = ''.join(random.choice(string.ascii_letters) for _ in range(32))
     ext = file.filename.split('.')[-1]
+    if ext not in SUPPORTED_FILE_EXTENSIONS:
+        raise HTTPException(status_code=422, detail="Unsupported file extension")
     path = os.path.join(settings.STATIC_PATH, "photo", "lecturer", f"{random_string}.{ext}")
     async with aiofiles.open(path, 'wb') as out_file:
         content = await file.read()

tests/conftest.py

@@ -101,6 +101,8 @@ def photo_path(client_auth: TestClient, dbsession: Session, lecturer_path: str):
     dbsession.commit()
 
 
+
+
 @pytest.fixture()
 def event_path(client_auth: TestClient, dbsession: Session, lecturer_path, room_path, group_path):
     RESOURCE = f"/timetable/event/"

tests/lecturer/photo.notpng

@@ -1,7 +1,9 @@
+import os
+
 from fastapi.testclient import TestClient
 from starlette import status
-from calendar_backend.settings import get_settings
 
+from calendar_backend.settings import get_settings
 
 settings = get_settings()
 settings.STATIC_PATH = './static'
@@ -25,3 +27,11 @@ def test_delete(client_auth: TestClient, photo_path: str):
 
     response = client_auth.get(photo_path)
     assert response.status_code == status.HTTP_404_NOT_FOUND
+
+
+def test_unsupported_format(lecturer_path: str, client_auth: TestClient):
+    RESOURCE = f"{lecturer_path}/photo"
+    with open(os.path.dirname(__file__) + "/photo.notpng", "rb") as f:
+        response = client_auth.post(RESOURCE, files={"photo": f})
+    assert response.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY
+