diff --git a/database/__init__.py b/database/__init__.py index 4ff63c3..38566ee 100644 --- a/database/__init__.py +++ b/database/__init__.py @@ -2,11 +2,12 @@ from models import settings from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker from sqlalchemy.orm import DeclarativeBase -from .user import * - engine = create_async_engine(settings.database) sessions = async_sessionmaker(engine) class Base(DeclarativeBase): pass + + +from .user import User diff --git a/main.py b/main.py index 9161668..e6f6013 100644 --- a/main.py +++ b/main.py @@ -18,6 +18,7 @@ async def lifespan(app: FastAPI): await connection.run_sync(database.Base.metadata.create_all) yield + app = FastAPI(lifespan=lifespan) app.include_router(routes.router) diff --git a/models/__init__.py b/models/__init__.py index 030e7b4..93ff70c 100644 --- a/models/__init__.py +++ b/models/__init__.py @@ -3,7 +3,7 @@ import pydantic class BaseModel(pydantic.BaseModel): class Config: - orm_mode = True + from_attributes = True from .settings import settings diff --git a/models/user.py b/models/user.py index ff202d3..3281f32 100644 --- a/models/user.py +++ b/models/user.py @@ -17,3 +17,11 @@ class Token(BaseModel): id: int username: str token: str + + +class DeleteUser(BaseModel): + username: str + + +class UpdatePassword(BaseModel): + password: str diff --git a/poetry.lock b/poetry.lock index 0afc738..a455c59 100644 --- a/poetry.lock +++ b/poetry.lock @@ -533,6 +533,23 @@ azure-key-vault = ["azure-identity (>=1.16.0)", "azure-keyvault-secrets (>=4.8.0 toml = ["tomli (>=2.0.1)"] yaml = ["pyyaml (>=6.0.1)"] +[[package]] +name = "pyjwt" +version = "2.9.0" +description = "JSON Web Token implementation in Python" +optional = false +python-versions = ">=3.8" +files = [ + {file = "PyJWT-2.9.0-py3-none-any.whl", hash = "sha256:3b02fb0f44517787776cf48f2ae25d8e14f300e6d7545a4315cee571a415e850"}, + {file = "pyjwt-2.9.0.tar.gz", hash = "sha256:7e1e5b56cc735432a7369cbfa0efe50fa113ebecdc04ae6922deba8b84582d0c"}, +] + +[package.extras] +crypto = ["cryptography (>=3.4.0)"] +dev = ["coverage[toml] (==5.0.4)", "cryptography (>=3.4.0)", "pre-commit", "pytest (>=6.0.0,<7.0.0)", "sphinx", "sphinx-rtd-theme", "zope.interface"] +docs = ["sphinx", "sphinx-rtd-theme", "zope.interface"] +tests = ["coverage[toml] (==5.0.4)", "pytest (>=6.0.0,<7.0.0)"] + [[package]] name = "python-dotenv" version = "1.0.1" @@ -694,4 +711,4 @@ standard = ["colorama (>=0.4)", "httptools (>=0.5.0)", "python-dotenv (>=0.13)", [metadata] lock-version = "2.0" python-versions = "^3.11" -content-hash = "c63b5753bf55ec6d3ae3cf225ad9d388fa1705cc2ce49718822221a8c19b2d66" +content-hash = "1d617553a03b53b743a89007a49b7167e21a07b5c19277b56b8b436af92c1f00" diff --git a/pyproject.toml b/pyproject.toml index 59d4c4c..bdca83a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,6 +14,7 @@ pydantic = "^2.8.2" SQLAlchemy = "^2.0.32" pydantic-settings = "^2.4.0" asyncpg = "^0.29.0" +pyjwt = "^2.9.0" [tool.poetry.group.dev.dependencies] isort = "^5.13.2" diff --git a/requirements.txt b/requirements.txt index 43a8671..4a5088f 100644 Binary files a/requirements.txt and b/requirements.txt differ diff --git a/routes/__init__.py b/routes/__init__.py index 89e1dda..246f78b 100644 --- a/routes/__init__.py +++ b/routes/__init__.py @@ -1,7 +1,9 @@ from fastapi import APIRouter from . import admin +from . import user router = APIRouter() router.include_router(admin.router) +router.include_router(user.router) diff --git a/routes/admin.py b/routes/admin.py index 08cdf94..69423c2 100644 --- a/routes/admin.py +++ b/routes/admin.py @@ -1,29 +1,15 @@ import secrets -from typing import Annotated -import hashlib -from fastapi import APIRouter, Depends, HTTPException -from sqlalchemy import select +from fastapi import APIRouter, HTTPException +from sqlalchemy import select, delete import database -from models import settings, user +from models import settings, user, DeleteUser +from .utils import Admin, hash_password router = APIRouter(prefix="/admin") -def hash_password(password: str, salt: str) -> str: - return hashlib.sha512((password + salt).encode('utf-8')).hexdigest() - - -def verify_admin(token: str): - if token != settings.admin_password: - raise HTTPException(401, "Unauthorized") - return True - - -Admin = Annotated[bool, Depends(verify_admin, use_cache=False)] - - @router.post("/user") async def create_user(auth: user.Auth, admin_token: Admin): if len(auth.username.strip()) == 0: @@ -36,17 +22,26 @@ async def create_user(auth: user.Auth, admin_token: Admin): salt = secrets.token_hex(8) async with database.sessions.begin() as session: - stmt = select(database.User).where(database.User.username == auth.username) - db_user = session.execute(stmt).scalar_one_or_none() - if db_user is not None: + stmt = select(database.User).where( + database.User.username == auth.username.strip() + ) + db_request = await session.execute(stmt) + user = db_request.scalar_one_or_none() + if user is not None: raise HTTPException(400, "User with this username already exists") new_user = database.User( - username=auth.username, - password=hash_password(auth.password, salt), + username=auth.username.strip(), + password=hash_password(auth.password.strip(), salt), salt=salt, ) session.add(new_user) - await session.flush() - return {'status': 'Success'} \ No newline at end of file + +@router.delete("/user") +async def delete_user(user: DeleteUser, admin_token: Admin): + async with database.sessions.begin() as session: + stmt = delete(database.User).where( + database.User.username == user.username.strip() + ) + await session.execute(stmt) diff --git a/routes/user.py b/routes/user.py index 5014484..be346a7 100644 --- a/routes/user.py +++ b/routes/user.py @@ -1,8 +1,47 @@ -import hashlib +import secrets -from fastapi import APIRouter -from fastapi.security import OAuth2PasswordBearer +import jwt +from fastapi import APIRouter, HTTPException +from sqlalchemy import select import database +import models +from models import settings +from .utils import hash_password, User router = APIRouter(prefix="/user") + + +@router.post("/login") +async def login(auth: models.Auth): + async with database.sessions.begin() as session: + stmt = select(database.User).where( + database.User.username == auth.username.strip() + ) + request = await session.execute(stmt) + user = request.scalar_one_or_none() + + if ( + user is None + or hash_password(auth.password.strip(), user.salt) != user.password + ): + raise HTTPException(403, "Forbidden") + + return models.Token( + id=user.id, + username=user.username, + token=jwt.encode( + {"sub": user.id}, settings.secret + user.password, "HS256" + ), + ) + + +@router.put("/update/password") +async def update_password(user: User, new: models.UpdatePassword): + if len(new.password.strip()) == 0: + raise HTTPException(400, "Password must not be empty") + + async with database.sessions.begin() as session: + session.add(user) + user.salt = secrets.token_hex(8) + user.password = hash_password(new.password.strip(), user.salt) diff --git a/routes/utils.py b/routes/utils.py new file mode 100644 index 0000000..b147b9f --- /dev/null +++ b/routes/utils.py @@ -0,0 +1,50 @@ +from typing import Annotated +import hashlib + +import jwt +from fastapi import Depends, HTTPException, Header +from sqlalchemy import select + +from models import settings +import database + + +def hash_password(password: str, salt: str) -> str: + return hashlib.sha512((password + salt).encode("utf-8")).hexdigest() + + +def verify_admin(password: Annotated[str, Header(alias="x-token")]): + if password.strip() != settings.admin_password: + raise HTTPException(401, "Unauthorized") + return True + + +async def verify_user(token: Annotated[str, Header(alias="x-token")]) -> database.User: + try: + data = jwt.decode( + token, algorithms=["HS256"], options={"verify_signature": False} + ) + except jwt.exceptions.DecodeError: + raise HTTPException(401, "Invalid token") + + if "sub" not in data and not isinstance(data["sub"], int): + raise HTTPException(401, "Invalid token") + + async with database.sessions.begin() as session: + stmt = select(database.User).where(database.User.id == data.get("sub")) + db_request = await session.execute(stmt) + user = db_request.scalar_one_or_none() + + if user is None: + raise HTTPException(401, "Invalid token") + + try: + jwt.decode(token, settings.secret + user.password, algorithms=["HS256"]) + except jwt.exceptions.InvalidSignatureError: + raise HTTPException(401, "Invalid token") + + return user + + +User = Annotated[database.User, Depends(verify_user, use_cache=False)] +Admin = Annotated[bool, Depends(verify_admin, use_cache=False)]