nas-tools/dashboard/backend/auth.py

from datetime import datetime, timedelta, timezone
from typing import Optional
import time
import threading
import tempfile
import jwt
from passlib.context import CryptContext
import pyotp
from fastapi import Depends, HTTPException, Response, 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")

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: Optional[timedelta] = None):
    to_encode = data.copy()
    expire = datetime.now(timezone.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(timezone.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
AUTH_FILE = config.VOLUME_ROOT + "/docker/nas-dashboard/auth.json"
import json, os
import base64
import hashlib
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(AUTH_FILE, data)
        return result

def _get_fernet():
    from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
    from cryptography.hazmat.primitives import hashes
    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:
            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):
    _atomic_json_write(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):
        return False

    def mutator(data):
        current_ts = int(time.time()) // 30
        if data.get("last_totp_ts") == current_ts:
            return False
        data["last_totp_ts"] = current_ts
        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()