Merge pull request #142 from Delemangi/feature/pass-encryption

Feature/pass encryption

.vscode/settings.json

@@ -70,5 +70,6 @@
   },
   "svelte.enable-ts-plugin": true,
   "svelte.plugin.svelte.format.config.singleQuote": true,
-  "svelte.plugin.svelte.format.config.svelteStrictMode": true
+  "svelte.plugin.svelte.format.config.svelteStrictMode": true,
+  "vitest.disableWorkspaceWarning": true
 }

README.md

@@ -23,6 +23,7 @@ Authors:
 - File previews
 - File whitelist
 - Encryption
+- 2FA
 - User registration and login
 - Intuitive interface
 - Support for webhooks

backend/alembic/versions/0f18e18f9ae9_add_2fa_auth.py

@@ -0,0 +1,30 @@
+"""Add 2FA auth
+
+Revision ID: 0f18e18f9ae9
+Revises: adbfe149f408
+Create Date: 2024-06-11 00:00:10.056063
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = '0f18e18f9ae9'
+down_revision: Union[str, None] = 'adbfe149f408'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column('user', sa.Column('code_2fa', sa.String(), nullable=True))
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column('user', 'code_2fa')
+    # ### end Alembic commands ###

backend/app/auth/exceptions.py

@@ -15,3 +15,7 @@
 USERNAME_TAKEN_EXCEPTION = HTTPException(
     status_code=status.HTTP_406_NOT_ACCEPTABLE, detail="Username already taken."
 )
+
+AUTHENTICATION_2FA_EXCEPTION = HTTPException(
+    status_code=status.HTTP_302_FOUND, detail="User found. 2FA authentication required."
+)

backend/app/auth/models.py

@@ -27,6 +27,7 @@ class User(Base):
         "Webhook", back_populates="user", cascade="all, delete-orphan", lazy="selectin"
     )
     shared_files = relationship("Share", back_populates="user", lazy="selectin")
+    code_2fa = Column(String, nullable=True)
 
     def has_remaining_quota(self: Self) -> bool:
         return bool(self.quota != 0)

backend/app/auth/router.py

@@ -1,19 +1,24 @@
 from typing import Annotated
 
-from fastapi import APIRouter, Depends
+from fastapi import APIRouter, Depends, Form
 from fastapi.security import OAuth2PasswordRequestForm
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from ..database import get_async_session
 from ..schemas import RequestStatus
-from .exceptions import CREDENTIALS_EXCEPTION
-from .schemas import Token, User
+from .dependencies import get_current_user
+from .exceptions import AUTHENTICATION_2FA_EXCEPTION, CREDENTIALS_EXCEPTION
+from .models import User as DbUser
+from .schemas import Code2FA, Token, User, UserMetadata
 from .service import (
     authenticate_user,
     create_access_token,
     create_user,
     oauth2_scheme,
+    remove_2fa_code,
     remove_token,
+    update_2fa_code,
+    verify_2fa_code,
 )
 
 router = APIRouter(tags=["auth"])
@@ -35,12 +40,16 @@ async def test(session: Annotated[AsyncSession, Depends(get_async_session)]) ->
 async def login_for_access_token(
     form_data: Annotated[OAuth2PasswordRequestForm, Depends()],
     session: Annotated[AsyncSession, Depends(get_async_session)],
+    code_2fa: str = Form(None),
 ) -> dict[str, str]:
     user = await authenticate_user(form_data.username, form_data.password, session)
 
     if not user:
         raise CREDENTIALS_EXCEPTION
 
+    if not verify_2fa_code(user, code_2fa):
+        raise AUTHENTICATION_2FA_EXCEPTION
+
     access_token = await create_access_token(session, data={"sub": str(user.username)})
     return {"access_token": access_token, "token_type": "bearer"}
 
@@ -60,3 +69,44 @@ async def register(
 ) -> RequestStatus:
     user = await create_user(user_schema.username, user_schema.password, 30, session)
     return RequestStatus(message=f"User {user.username} registered successfully")
+
+
+@router.post("/2fa")
+async def get_2fa_token(
+    form_data: Annotated[OAuth2PasswordRequestForm, Depends()],
+    session: Annotated[AsyncSession, Depends(get_async_session)],
+) -> Code2FA:
+    user = await authenticate_user(form_data.username, form_data.password, session)
+
+    if not user:
+        raise CREDENTIALS_EXCEPTION
+
+    code_2fa = await update_2fa_code(user, session)
+
+    return Code2FA(code=code_2fa)
+
+
+@router.post("/2fa/disable")
+async def disable_2fa(
+    form_data: Annotated[OAuth2PasswordRequestForm, Depends()],
+    session: Annotated[AsyncSession, Depends(get_async_session)],
+) -> RequestStatus:
+    user = await authenticate_user(form_data.username, form_data.password, session)
+
+    if not user:
+        raise CREDENTIALS_EXCEPTION
+
+    await remove_2fa_code(user, session)
+
+    return RequestStatus(message="2FA disabled successfully")
+
+
+@router.get("/fetch_user_data", response_model=str)
+async def fetch_user_data(
+    current_user: Annotated[DbUser, Depends(get_current_user)],
+) -> UserMetadata:
+    return UserMetadata(
+        username=str(current_user.username),
+        quota=int(current_user.quota.value),
+        is_2fa_enabled=(current_user.code_2fa is not None),
+    )

backend/app/auth/schemas.py

@@ -17,3 +17,13 @@ class Token(BaseModel):
 
 class TokenData(BaseModel):
     username: str | None = None
+
+
+class Code2FA(BaseModel):
+    code: str
+
+
+class UserMetadata(BaseModel):
+    username: str
+    quota: int
+    is_2fa_enabled: bool

backend/app/auth/service.py

@@ -4,10 +4,11 @@
 from datetime import UTC, datetime, timedelta
 from uuid import UUID
 
+import pyotp
 from fastapi.security import OAuth2PasswordBearer
 from jose import jwt
 from passlib.context import CryptContext
-from sqlalchemy import ColumnElement, and_, delete, select
+from sqlalchemy import ColumnElement, and_, delete, select, update
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from ..database import get_async_session
@@ -58,6 +59,29 @@ async def authenticate_user(username: str, password: str, session: AsyncSession)
     )
 
 
+def verify_2fa_code(user: User, code: str | None) -> bool:
+    if user.code_2fa is None and code is None:
+        return True
+    if user.code_2fa is None or code is None:
+        return False
+
+    totp = pyotp.TOTP(str(user.code_2fa))
+    return totp.verify(code)
+
+
+async def update_2fa_code(user: User, session: AsyncSession) -> str:
+    code = pyotp.random_base32()
+    await session.execute(update(User).where(User.id == user.id).values(code_2fa=code))
+    await session.commit()
+
+    return code
+
+
+async def remove_2fa_code(user: User, session: AsyncSession) -> None:
+    await session.execute(update(User).where(User.id == user.id).values(code_2fa=None))
+    await session.commit()
+
+
 async def create_access_token(
     session: AsyncSession,
     data: dict[str, str | int | datetime],

backend/app/files/router.py

@@ -1,9 +1,10 @@
+import io
 from datetime import datetime, timedelta
 from pathlib import Path
 from typing import Annotated
 
-from fastapi import APIRouter, Depends, Form, UploadFile
-from fastapi.responses import FileResponse
+from fastapi import APIRouter, Depends, Form, Header, UploadFile
+from fastapi.responses import StreamingResponse
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from ..auth.dependencies import get_current_user
@@ -19,9 +20,9 @@
     delete_file,
     get_all_files_user,
     get_metadata_path,
-    upload_file_unencrypted,
+    upload_file,
+    verify_and_decrypt_file,
     verify_file,
-    verify_file_link,
 )
 
 router = APIRouter(tags=["files"])
@@ -68,9 +69,10 @@ async def create_upload_file(
     current_user: Annotated[User, Depends(get_current_user)],
     session: Annotated[AsyncSession, Depends(get_async_session)],
     file: UploadFile,
+    password: str = Form(""),
     is_shared: bool = Form(False),
 ) -> FileUploaded:
-    new_file = await upload_file_unencrypted(session, file, current_user, is_shared)
+    new_file = await upload_file(session, file, current_user, is_shared, password)
     return FileUploaded(filename=new_file, username=str(current_user.username))
 
 
@@ -87,19 +89,16 @@ async def get_file(
     path: str,
     current_user: Annotated[User, Depends(get_current_user)],
     session: Annotated[AsyncSession, Depends(get_async_session)],
-) -> FileResponse:
-    filename = await verify_file(path, current_user, session)
-    return FileResponse(FILE_PATH + path, filename=filename)
-
-
-@router.get("/download-link/{path}")
-async def get_file_link(
-    path: str,
-    session: Annotated[AsyncSession, Depends(get_async_session)],
-    current_user: Annotated[User, Depends(get_current_user)],
-) -> FileResponse:
-    filename = await verify_file_link(path, session, current_user)
-    return FileResponse(FILE_PATH + path, filename=filename)
+    password: str = Header(None),
+) -> StreamingResponse:
+    print(path)
+    file = await verify_and_decrypt_file(path, current_user, session, password)
+    file_object = io.BytesIO(file)
+    return StreamingResponse(
+        file_object,
+        media_type="application/octet-stream",
+        headers={"Content-Disposition": f"attachment; filename={path}"},
+    )
 
 
 @router.get("/metadata/{path}", response_model=MetadataFileResponse)