nas-tools/dashboard/backend/auth_service.py

import tempfile
import threading
import time
from datetime import UTC, datetime, timedelta

import jwt
import pyotp
from fastapi import Depends, HTTPException, Response, status
from fastapi.security import OAuth2PasswordBearer
from passlib.context import CryptContext

import config

# Password handling
pwd_context = CryptContext(schemes=["pbkdf2_sha256"], deprecated="auto")
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="api/auth/login")

ACCESS_COOKIE_NAME = "nas_access_token"
REFRESH_COOKIE_NAME = "nas_refresh_token"
COOKIE_SECURE = True
COOKIE_SAMESITE = "lax"
COOKIE_PATH = "/"


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: timedelta | None = None):
    to_encode = data.copy()
    expire = datetime.now(UTC) + (expires_delta or timedelta(minutes=15))
    to_encode.update({"exp": expire, "type": "access"})
    return jwt.encode(to_encode, config.SECRET_KEY, algorithm=config.ALGORITHM)


def create_refresh_token(data: dict):
    to_encode = data.copy()
    expire = datetime.now(UTC) + timedelta(minutes=config.REFRESH_TOKEN_EXPIRE_MINUTES)
    to_encode.update({"exp": expire, "type": "refresh", "tv": _load_token_version()})
    return jwt.encode(to_encode, config.SECRET_KEY, algorithm=config.ALGORITHM)


def _cookie_max_age(minutes: int) -> int:
    return max(60, int(minutes * 60))


def set_auth_cookies(response: Response, access_token: str, refresh_token: str):
    response.set_cookie(
        key=ACCESS_COOKIE_NAME,
        value=access_token,
        httponly=True,
        secure=COOKIE_SECURE,
        samesite=COOKIE_SAMESITE,
        path=COOKIE_PATH,
        max_age=_cookie_max_age(config.ACCESS_TOKEN_EXPIRE_MINUTES),
    )
    response.set_cookie(
        key=REFRESH_COOKIE_NAME,
        value=refresh_token,
        httponly=True,
        secure=COOKIE_SECURE,
        samesite=COOKIE_SAMESITE,
        path=COOKIE_PATH,
        max_age=_cookie_max_age(config.REFRESH_TOKEN_EXPIRE_MINUTES),
    )


def clear_auth_cookies(response: Response):
    response.delete_cookie(ACCESS_COOKIE_NAME, path=COOKIE_PATH)
    response.delete_cookie(REFRESH_COOKIE_NAME, path=COOKIE_PATH)


def user_from_access_token(token: str, include_auth_header: bool = True):
    from rbac import User, get_pages, get_sidebar_links

    credentials_exception = HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="Could not validate credentials",
        headers={"WWW-Authenticate": "Bearer"} if include_auth_header else None,
    )
    try:
        payload = jwt.decode(token, config.SECRET_KEY, algorithms=[config.ALGORITHM])
        if payload.get("type") != "access":
            raise credentials_exception
        username: str = payload.get("sub")
        role: str = payload.get("role", "admin")
        if username is None:
            raise credentials_exception
    except jwt.PyJWTError:
        raise credentials_exception

    pages = get_pages(username, role)
    sidebar_links = get_sidebar_links(username, role)
    return User(username=username, role=role, pages=pages, sidebar_links=sidebar_links)


async def get_current_user(token: str = Depends(oauth2_scheme)):
    return user_from_access_token(token)


async def get_current_user_ws(token: str):
    return user_from_access_token(token, include_auth_header=False)


# TOTP Persistence
def get_auth_file():
    return config.VOLUME_ROOT + "/docker/nas-dashboard/auth.json"


import base64
import hashlib
import json
import os

from cryptography.fernet import Fernet

_auth_lock = threading.RLock()


def _atomic_json_write(path: str, data: dict):
    os.makedirs(os.path.dirname(path), exist_ok=True)
    fd, temp_path = tempfile.mkstemp(dir=os.path.dirname(path), prefix=".tmp-auth-", suffix=".json")
    try:
        with os.fdopen(fd, "w") as f:
            json.dump(data, f)
            f.flush()
            os.fsync(f.fileno())
        os.replace(temp_path, path)
    finally:
        if os.path.exists(temp_path):
            os.unlink(temp_path)


def _update_auth_data(mutator):
    with _auth_lock:
        data = _load_auth_data()
        result = mutator(data)
        _atomic_json_write(get_auth_file(), data)
        return result


def _get_fernet():
    from cryptography.hazmat.primitives import hashes
    from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC

    kdf = PBKDF2HMAC(algorithm=hashes.SHA256(), length=32, salt=b"nas-dashboard-fernet", iterations=480000)
    key = base64.urlsafe_b64encode(kdf.derive(config.SECRET_KEY.encode()))
    return Fernet(key)


def _get_legacy_fernet():
    key = base64.urlsafe_b64encode(hashlib.sha256(config.SECRET_KEY.encode()).digest())
    return Fernet(key)


def _encrypt(plaintext: str) -> str:
    if not plaintext:
        return ""
    return _get_fernet().encrypt(plaintext.encode()).decode()


def _decrypt(ciphertext: str) -> str:
    if not ciphertext:
        return ""
    try:
        return _get_fernet().decrypt(ciphertext.encode()).decode()
    except Exception:
        pass
    # Fallback: try legacy SHA256-based key
    try:
        return _get_legacy_fernet().decrypt(ciphertext.encode()).decode()
    except Exception:
        return ciphertext  # fallback for unencrypted legacy values


def _load_auth_data():
    with _auth_lock:
        try:
            auth_path = get_auth_file()
            if os.path.exists(auth_path):
                with open(auth_path) as f:
                    return json.load(f)
        except Exception:
            pass
        return {}


def _save_auth_data(data):
    _atomic_json_write(get_auth_file(), data)


def load_totp_secret():
    encrypted = _load_auth_data().get("totp_secret", "")
    if not encrypted:
        return config.TOTP_SECRET
    return _decrypt(encrypted)


def save_totp_secret(secret: str):
    def mutator(data):
        data["totp_secret"] = _encrypt(secret) if secret else ""

    _update_auth_data(mutator)


def load_password_hash():
    return _load_auth_data().get("password_hash", "") or config.ADMIN_PASSWORD_HASH


def save_password_hash(hashed: str):
    def mutator(data):
        data["password_hash"] = hashed

    _update_auth_data(mutator)


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)
    if not totp.verify(token, valid_window=1):  # Allow 1 time step before/after (±30s)
        return False

    def mutator(data):
        # Track recently used codes to prevent replay attacks
        used_codes = data.get("used_totp_codes", [])
        current_time = int(time.time())

        # Check if this code was already used
        if token in [code for code, _ in used_codes]:
            return False

        # Add current code with timestamp
        used_codes.append((token, current_time))

        # Clean up codes older than 90 seconds (3 time windows)
        used_codes = [(code, ts) for code, ts in used_codes if current_time - ts < 90]

        # Keep only last 5 codes to limit memory
        data["used_totp_codes"] = used_codes[-5:]
        return True

    return _update_auth_data(mutator)


# Passkey credentials
def load_passkey_credentials():
    return _load_auth_data().get("passkey_credentials", [])


def save_passkey_credential(cred):
    def mutator(data):
        creds = data.get("passkey_credentials", [])
        creds.append(cred)
        data["passkey_credentials"] = creds

    _update_auth_data(mutator)


def clear_passkey_credentials():
    def mutator(data):
        data["passkey_credentials"] = []

    _update_auth_data(mutator)


# Token version for refresh token rotation
def _load_token_version() -> int:
    return _load_auth_data().get("token_version", 0)


def increment_token_version() -> int:
    def mutator(data):
        v = data.get("token_version", 0) + 1
        data["token_version"] = v
        return v

    return _update_auth_data(mutator)


def verify_token_version(tv: int) -> bool:
    return tv == _load_token_version()