Spaces:
Sleeping
Sleeping
File size: 3,364 Bytes
8b1e853 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 |
from fastapi import APIRouter, Depends, HTTPException, status
from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
from jose import JWTError, jwt
from passlib.context import CryptContext
from pydantic import BaseModel
from datetime import datetime, timedelta, timezone
import sqlite3
import os
SECRET_KEY = "your-secret-key"
ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_MINUTES = 300
pwd_context = CryptContext(schemes=["django_pbkdf2_sha256"], deprecated="auto")
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
auth_router = APIRouter()
class User(BaseModel):
username: str
class Token(BaseModel):
access_token: str
token_type: str
# SQLite setup
DB_NAME = "users.db"
def init_db():
conn = sqlite3.connect(DB_NAME)
cur = conn.cursor()
cur.execute('''CREATE TABLE IF NOT EXISTS users
(username TEXT PRIMARY KEY, hashed_password TEXT)''')
# Add sample users if they don't exist
cur.execute("INSERT OR IGNORE INTO users VALUES (?, ?)",
("admin", pwd_context.hash("admin")))
conn.commit()
conn.close()
# Initialize the database
init_db()
def verify_password(plain_password, hashed_password):
return pwd_context.verify(plain_password, hashed_password)
def get_user(username: str):
conn = sqlite3.connect(DB_NAME)
cur = conn.cursor()
cur.execute("SELECT * FROM users WHERE username=?", (username,))
user = cur.fetchone()
conn.close()
if user:
return User(username=user[0])
return None
def authenticate_user(username: str, password: str):
conn = sqlite3.connect(DB_NAME)
cur = conn.cursor()
cur.execute("SELECT * FROM users WHERE username=?", (username,))
user = cur.fetchone()
conn.close()
if not user:
return False
if not verify_password(password, user[1]):
return False
return User(username=user[0])
def create_access_token(data: dict):
to_encode = data.copy()
expire = datetime.now(timezone.utc) + timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
to_encode.update({"exp": expire})
encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
return encoded_jwt
@auth_router.post("/token", response_model=Token)
async def login(form_data: OAuth2PasswordRequestForm = Depends()):
user = authenticate_user(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 = create_access_token(data={"sub": user.username})
return {"access_token": access_token, "token_type": "bearer"}
async def get_current_user(token: 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: str = payload.get("sub")
if username is None:
raise credentials_exception
user = get_user(username)
if user is None:
raise credentials_exception
return user
except JWTError:
raise credentials_exception |