feat: use jti to revoke token (#30)

Author: msterhuj

fob_api/__init__.py

@@ -1,5 +1,5 @@
 from .config import Config
-from .database import init_engine
+from .database import init_engine, get_session
 from .lib.headscale import HeadScale
 from .vpn import headscale_driver
 from . import mail

fob_api/auth/__init__.py

@@ -12,6 +12,7 @@
 
 from fob_api.config import Config
 from fob_api.models.database import User
+from fob_api.models.database import Token as TokenDB
 from fob_api import engine
 
 password_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
@@ -20,7 +21,7 @@
 
 oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/token")
 jwt_secret = Config().jwt_secret_key
-jwt_expire_days = 1
+jwt_expire_days = 15
 
 if not jwt_secret:
     raise ValueError("JWT secret not set")
@@ -33,29 +34,30 @@ def hash_password(password: str) -> str:
     """
     return password_context.hash(password)
 
-
-def encode_token(username) -> str:
-    return jwt.encode({
+def make_token_data(username: str) -> dict:
+    return {
         "exp": datetime.now() + timedelta(days=jwt_expire_days),
         "iat": datetime.now(),
         "jti": str(uuid4()),
         "nbf": datetime.now(),
         "sub": str(username)
-    }, jwt_secret, algorithm="HS256")
+    }
+
+def encode_token(token_data) -> str:
+    return jwt.encode(token_data, jwt_secret, algorithm="HS256")
 
 
 def get_current_user(token: Annotated[str, Depends(oauth2_scheme)]) -> User:
     try:
         payload = jwt.decode(token, jwt_secret, algorithms=["HS256"])
         with Session(engine) as session:
-            return session.exec(select(User).where(User.username == payload["sub"])).first()
-        raise HTTPException(status_code=401, detail="No user matching the token")
-    except JWTClaimsError as e:
-        raise HTTPException(status_code=401, detail=f"JWTClaimsError: {e}")
-    except ExpiredSignatureError as e:
-        raise HTTPException(status_code=401, detail=f"ExpiredSignatureError: {e}")
-    except JWTError as e:
-        raise HTTPException(status_code=401, detail=f"JWTError: {e}")
+            user = session.exec(select(User).where(User.username == payload["sub"])).first()
+            token = session.exec(select(TokenDB).where(TokenDB.token_id == payload["jti"])).first()
+            if user and token:
+                return user
+        raise HTTPException(status_code=401, detail="Invalid token")
+    except (JWTClaimsError, ExpiredSignatureError, JWTError) as e:
+        raise HTTPException(status_code=401, detail=f"JWT Error: {e}")
 
 
 def basic_auth_validator(username: str, password: str) -> User:

fob_api/database.py

@@ -1,3 +1,5 @@
+from fastapi import Depends, FastAPI, HTTPException, Query
+from sqlmodel import Field, Session, SQLModel, create_engine, select
 from sqlalchemy import Engine
 from sqlmodel import create_engine, SQLModel
 from fob_api import Config
@@ -9,3 +11,9 @@ def init_engine() -> Engine:
     """
     print("Initializing database engine")
     return create_engine(Config().database_url, echo=False, pool_recycle=1800, pool_pre_ping=True)
+
+engine = init_engine()
+
+def get_session():
+    with Session(engine) as session:
+        yield session

fob_api/models/database/__init__.py

@@ -1,6 +1,7 @@
 from .user import (
     User,
-    UserPasswordReset
+    UserPasswordReset,
+    Token
 )
 from .headscale import (
     HeadScalePolicyACL,

fob_api/models/database/user.py

@@ -24,3 +24,13 @@ class UserPasswordReset(SQLModel, table=True):
     source_ip: str
     created_at: datetime = Field(default=datetime.now())
     expires_at: datetime
+
+class Token(SQLModel, table=True):
+    """
+    This class represents the Token
+    """
+    id: int = Field(primary_key=True)
+    expires_at: datetime
+    created_at: datetime
+    token_id: str
+    user_id: int = Field(foreign_key="user.id")

fob_api/routes/token.py

@@ -2,27 +2,55 @@
 
 from fastapi import APIRouter, Depends, HTTPException
 from fastapi.security import OAuth2PasswordRequestForm
+from sqlmodel import Session, select
 
-from fob_api import auth
+from fob_api import auth, get_session
 from fob_api.models.database import User
+from fob_api.models.database import Token as TokenDB
 from fob_api.models.api import Token, TokenValidate
 
 router = APIRouter()
 
 @router.post("/token", response_model=Token, tags=["token"])
-def get_token(form_data: Annotated[OAuth2PasswordRequestForm, Depends()]) -> str:
+def get_token(form_data: Annotated[OAuth2PasswordRequestForm, Depends()], session: Session = Depends(get_session)) -> Token:
     user = auth.basic_auth_validator(form_data.username, form_data.password)
     if not user:
         raise HTTPException(status_code=401, detail="Invalid credentials")
-    token = auth.encode_token(user.username)
+    token_data = auth.make_token_data(user.username)
+    token_db: TokenDB = TokenDB(
+        expires_at=token_data["exp"],
+        created_at=token_data["iat"],
+        token_id=token_data["jti"],
+        user_id=user.id,
+    )
+    session.add(token_db)
+    session.commit()
+    token = auth.encode_token(token_data)
     return Token(access_token=token, token_type="bearer")
 
 
 @router.get("/token/refreshtoken", response_model=Token, tags=["token"])
-def refresh_token(user: Annotated[User, Depends(auth.get_current_user)]) -> str:
-    token = auth.encode_token(user.username)
+def refresh_token(user: Annotated[User, Depends(auth.get_current_user)], session: Session = Depends(get_session)) -> Token:
+    token_data = auth.make_token_data(user.username)
+    token_db: TokenDB = TokenDB(
+        expires_at=token_data["exp"],
+        created_at=token_data["iat"],
+        token_id=token_data["jti"],
+        user_id=user.id,
+    )
+    session.add(token_db)
+    session.commit()
+    token = auth.encode_token(token_data)
     return Token(access_token=token, token_type="bearer")
 
+@router.delete("/token/{jti}", tags=["token"])
+def revoke_token(jti: str, user: Annotated[User, Depends(auth.get_current_user)], session: Session = Depends(get_session)) -> None:
+    token = session.exec(select(TokenDB).where(TokenDB.token_id == jti)).first()
+    if not token or token.user_id != user.id:
+        raise HTTPException(status_code=404, detail="Cant revoke token")
+    session.delete(token)
+    session.commit()
+
 @router.get("/token/verify", response_model=TokenValidate, tags=["token"])
-def verify_token(user: Annotated[User, Depends(auth.get_current_user)]) -> str:
+def verify_token(user: Annotated[User, Depends(auth.get_current_user)]) -> TokenValidate:
     return TokenValidate(valid=True)

migrations/versions/2025_02_05_1931-22f19d8927c2_add_token_register.py

@@ -0,0 +1,39 @@
+"""add token register
+
+Revision ID: 22f19d8927c2
+Revises: aa1757ace187
+Create Date: 2025-02-05 19:31:37.252993
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+import sqlmodel
+
+
+# revision identifiers, used by Alembic.
+revision: str = '22f19d8927c2'
+down_revision: Union[str, None] = 'aa1757ace187'
+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.create_table('token',
+    sa.Column('id', sa.Integer(), nullable=False),
+    sa.Column('expires_at', sa.DateTime(), nullable=False),
+    sa.Column('created_at', sa.DateTime(), nullable=False),
+    sa.Column('token_id', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+    sa.Column('user_id', sa.Integer(), nullable=False),
+    sa.ForeignKeyConstraint(['user_id'], ['user.id'], ),
+    sa.PrimaryKeyConstraint('id')
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table('token')
+    # ### end Alembic commands ###