require token to use API

main.py

@@ -1,8 +1,9 @@
 import sys
 import os
 import time
-from fastapi import FastAPI, UploadFile, File, HTTPException
+from fastapi import FastAPI, UploadFile, File, HTTPException, Depends, status
 from fastapi.responses import FileResponse
+from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
 import uvicorn
 import traceback
 import pickle
@@ -10,6 +11,12 @@ import shutil
 from pathlib import Path
 from contextlib import asynccontextmanager
 import pandas as pd
+from typing import Annotated
+from datetime import datetime, timedelta, timezone
+import jwt
+from jwt.exceptions import InvalidTokenError
+from passlib.context import CryptContext
+from pydantic import BaseModel
 
 current_dir = os.path.dirname(os.path.abspath(__file__))
 sys.path.append(os.path.join(current_dir, "meisai-check-ai"))
@@ -42,6 +49,105 @@ os.makedirs(os.path.join(current_dir, "data"), exist_ok=True)
 os.makedirs(os.path.join(current_dir, "uploads"), exist_ok=True)
 os.makedirs(os.path.join(current_dir, "outputs"), exist_ok=True)
 
+# Authentication related settings
+SECRET_KEY = "09d25e094faa6ca2556c818166b7a9563b93f7099f6f0f4caa6cf63b88e8d3e7"
+ALGORITHM = "HS256"
+ACCESS_TOKEN_EXPIRE_HOURS = 24  # Token expiration set to 24 hours
+
+# Password hashing context
+pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
+
+# OAuth2 scheme for token
+oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
+
+# User database models
+class Token(BaseModel):
+    access_token: str
+    token_type: str
+
+class TokenData(BaseModel):
+    username: str | None = None
+
+class User(BaseModel):
+    username: str
+    email: str | None = None
+    full_name: str | None = None
+    disabled: bool | None = None
+
+class UserInDB(User):
+    hashed_password: str
+
+# Fake users database with hashed passwords
+users_db = {
+    "chien_vm": {
+        "username": "chien_vm",
+        "full_name": "Chien VM",
+        "email": "chien_vm@detomo.co.jp",
+        "hashed_password": "$2b$12$RtcKFk7B3hKd7vYkwxdFN.eBXSiryQIRUG.OoJ07Pl9lzHNUkugMi",
+        "disabled": False,
+    },
+    "hoi_nv": {
+        "username": "hoi_nv",
+        "full_name": "Hoi NV",
+        "email": "hoi_nv@detomo.co.jp",
+        "hashed_password": "$2b$12$RtcKFk7B3hKd7vYkwxdFN.eBXSiryQIRUG.OoJ07Pl9lzHNUkugMi",
+        "disabled": False,
+    }
+}
+
+# Authentication helper functions
+def verify_password(plain_password, hashed_password):
+    return pwd_context.verify(plain_password, hashed_password)
+
+def get_user(db, username: str):
+    if username in db:
+        user_dict = db[username]
+        return UserInDB(**user_dict)
+    return None
+
+def authenticate_user(fake_db, username: str, password: str):
+    user = get_user(fake_db, username)
+    if not user:
+        return False
+    if not verify_password(password, user.hashed_password):
+        return False
+    return user
+
+def create_access_token(data: dict, expires_delta: timedelta | None = None):
+    to_encode = data.copy()
+    if expires_delta:
+        expire = datetime.now(timezone.utc) + expires_delta
+    else:
+        expire = datetime.now(timezone.utc) + timedelta(hours=ACCESS_TOKEN_EXPIRE_HOURS)
+    to_encode.update({"exp": expire})
+    encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
+    return encoded_jwt
+
+async def get_current_user(token: Annotated[str, Depends(oauth2_scheme)]):
+    credentials_exception = HTTPException(
+        status_code=status.HTTP_401_UNAUTHORIZED,
+        detail="Could not validate credentials",
+        headers={"WWW-Authenticate": "Bearer"},
+    )
+    try:
+        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
+        username = payload.get("sub")
+        if username is None:
+            raise credentials_exception
+        token_data = TokenData(username=username)
+    except InvalidTokenError:
+        raise credentials_exception
+    user = get_user(users_db, username=token_data.username)
+    if user is None:
+        raise credentials_exception
+    return user
+
+async def get_current_active_user(
+    current_user: Annotated[User, Depends(get_current_user)],
+):
+    if current_user.disabled:
+        raise HTTPException(status_code=400, detail="Inactive user")
+    return current_user
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
@@ -116,10 +222,34 @@ async def health_check():
     return {"status": "ok", "timestamp": time.time()}
 
 
+@app.post("/token")
+async def login_for_access_token(
+    form_data: Annotated[OAuth2PasswordRequestForm, Depends()]
+) -> Token:
+    """
+    Login endpoint to get an access token
+    """
+    user = authenticate_user(users_db, form_data.username, form_data.password)
+    if not user:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Incorrect username or password",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+    access_token_expires = timedelta(hours=ACCESS_TOKEN_EXPIRE_HOURS)
+    access_token = create_access_token(
+        data={"sub": user.username}, expires_delta=access_token_expires
+    )
+    return Token(access_token=access_token, token_type="bearer")
+
+
 @app.post("/predict")
-async def predict(file: UploadFile = File(...)):
+async def predict(
+    current_user: Annotated[User, Depends(get_current_active_user)],
+    file: UploadFile = File(...)
+):
     """
-    Process an input CSV file and return standardized names
+    Process an input CSV file and return standardized names (requires authentication)
     """
     global sentenceTransformerHelper, dic_standard_subject, sample_name_sentence_embeddings
     global sample_name_sentence_similarities, sampleData, name_groups
@@ -129,10 +259,10 @@ async def predict(file: UploadFile = File(...)):
 
     # Save uploaded file
     timestamp = int(time.time())
-    input_file_path = os.path.join(current_dir, "uploads", f"input_{timestamp}.csv")
+    input_file_path = os.path.join(current_dir, "uploads", f"input_{timestamp}_{current_user.username}.csv")
 
     # Use CSV format with correct extension
-    output_file_path = os.path.join(current_dir, "outputs", f"output_{timestamp}.csv")
+    output_file_path = os.path.join(current_dir, "outputs", f"output_{timestamp}_{current_user.username}.csv")
 
     try:
         with open(input_file_path, "wb") as buffer:
@@ -158,9 +288,7 @@ async def predict(file: UploadFile = File(...)):
         # Create output dataframe and save to CSV
         print("Columns of inputData.dataframe", inputData.dataframe.columns)
         inputData.dataframe.reset_index(drop=False, inplace=True)
-        columns_to_keep = ["ID", "シート名", "行", "科目", "名称", "摘要", "備考"]
-        output_df = inputData.dataframe[columns_to_keep].copy()
-        # Use .loc to avoid SettingWithCopyWarning
+        output_df = inputData.dataframe.copy()
         output_df.loc[:, "出力_科目"] = df_predicted["出力_科目"]
         output_df.loc[:, "出力_項目名"] = df_predicted["出力_項目名"]
         output_df.loc[:, "出力_確率度"] = df_predicted["出力_確率度"]
@@ -184,6 +312,6 @@ async def predict(file: UploadFile = File(...)):
         traceback.print_exc()
         raise HTTPException(status_code=500, detail=str(e))
 
-
+    
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=8000)