Gan, Jimmy
104 lines
3.3 KiB
Python
104 lines
3.3 KiB
Python
from datetime import datetime, timedelta
|
|
from typing import Optional
|
|
import jwt
|
|
from passlib.context import CryptContext
|
|
import pyotp
|
|
from fastapi import Depends, HTTPException, status
|
|
from fastapi.security import OAuth2PasswordBearer
|
|
import config
|
|
|
|
# Password handling
|
|
pwd_context = CryptContext(schemes=["pbkdf2_sha256"], deprecated="auto")
|
|
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="api/auth/login")
|
|
|
|
def verify_password(plain_password, hashed_password):
|
|
return pwd_context.verify(plain_password, hashed_password)
|
|
|
|
def get_password_hash(password):
|
|
return pwd_context.hash(password)
|
|
|
|
def create_access_token(data: dict, expires_delta: Optional[timedelta] = None):
|
|
to_encode = data.copy()
|
|
if expires_delta:
|
|
expire = datetime.utcnow() + expires_delta
|
|
else:
|
|
expire = datetime.utcnow() + timedelta(minutes=15)
|
|
to_encode.update({"exp": expire})
|
|
encoded_jwt = jwt.encode(to_encode, config.SECRET_KEY, algorithm=config.ALGORITHM)
|
|
return encoded_jwt
|
|
|
|
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, config.SECRET_KEY, algorithms=[config.ALGORITHM])
|
|
username: str = payload.get("sub")
|
|
if username is None:
|
|
raise credentials_exception
|
|
# In a real app, we'd check DB here.
|
|
# For this single-user system, we just check if it matches config.
|
|
if username != config.ADMIN_USER:
|
|
raise credentials_exception
|
|
except jwt.PyJWTError:
|
|
raise credentials_exception
|
|
return username
|
|
|
|
async def get_current_user_ws(token: str):
|
|
credentials_exception = HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail="Could not validate credentials",
|
|
)
|
|
try:
|
|
payload = jwt.decode(token, config.SECRET_KEY, algorithms=[config.ALGORITHM])
|
|
username: str = payload.get("sub")
|
|
if username is None or username != config.ADMIN_USER:
|
|
raise credentials_exception
|
|
except jwt.PyJWTError:
|
|
raise credentials_exception
|
|
return username
|
|
|
|
# TOTP Persistence
|
|
AUTH_FILE = config.VOLUME_ROOT + "/docker/nas-dashboard/auth.json"
|
|
import json, os
|
|
|
|
def _load_auth_data():
|
|
try:
|
|
if os.path.exists(AUTH_FILE):
|
|
with open(AUTH_FILE, "r") as f:
|
|
return json.load(f)
|
|
except Exception:
|
|
pass
|
|
return {}
|
|
|
|
def _save_auth_data(data):
|
|
os.makedirs(os.path.dirname(AUTH_FILE), exist_ok=True)
|
|
with open(AUTH_FILE, "w") as f:
|
|
json.dump(data, f)
|
|
|
|
def load_totp_secret():
|
|
return _load_auth_data().get("totp_secret", "") or config.TOTP_SECRET
|
|
|
|
def save_totp_secret(secret: str):
|
|
data = _load_auth_data()
|
|
data["totp_secret"] = secret
|
|
_save_auth_data(data)
|
|
|
|
def load_password_hash():
|
|
return _load_auth_data().get("password_hash", "") or config.ADMIN_PASSWORD_HASH
|
|
|
|
def save_password_hash(hashed: str):
|
|
data = _load_auth_data()
|
|
data["password_hash"] = hashed
|
|
_save_auth_data(data)
|
|
|
|
def verify_totp(token: str, secret: str = None):
|
|
if secret is None:
|
|
secret = load_totp_secret()
|
|
if not secret:
|
|
return True # 2FA not enabled
|
|
totp = pyotp.TOTP(secret)
|
|
return totp.verify(token)
|