nas-tools/dashboard/backend/routers/auth.py

from datetime import timedelta
import asyncio
import time
import json
import httpx
from fastapi import APIRouter, Depends, HTTPException, status, Request
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from slowapi import Limiter
from slowapi.util import get_remote_address
import jwt
import config
import auth
import pyotp
from webauthn import generate_registration_options, verify_registration_response, generate_authentication_options, verify_authentication_response
from webauthn.helpers.structs import PublicKeyCredentialDescriptor, UserVerificationRequirement
from webauthn.helpers import bytes_to_base64url, base64url_to_bytes, options_to_json

router = APIRouter()
limiter = Limiter(key_func=get_remote_address)

class LoginRequest(BaseModel):
    username: str
    password: str
    totp_code: str = ""

class Token(BaseModel):
    access_token: str
    refresh_token: str
    token_type: str
    user: str
    has_2fa: bool

class RefreshRequest(BaseModel):
    refresh_token: str

class Setup2FAResponse(BaseModel):
    secret: str
    uri: str

class Verify2FARequest(BaseModel):
    secret: str
    code: str

async def send_failed_login_alert(ip_address: str, username: str, reason: str):
    print(f"Triggering Telegram alert for {username} from {ip_address}: {reason}", flush=True)
    if not config.TELEGRAM_BOT_TOKEN or not config.TELEGRAM_CHAT_ID:
        print("Skipping Telegram alert: missing config", flush=True)
        return
    msg = f"🚨 *Dashboard Security Alert* 🚨\nFailed login attempt detected.\n\n👤 User: `{username}`\n🌐 IP: `{ip_address}`\n❌ Reason: {reason}"
    try:
        client_options = {}
        if config.TELEGRAM_PROXY:
            client_options["proxy"] = config.TELEGRAM_PROXY
        
        async with httpx.AsyncClient(**client_options) as client:
            res = await client.post(
                f"https://api.telegram.org/bot{config.TELEGRAM_BOT_TOKEN}/sendMessage",
                data={"chat_id": config.TELEGRAM_CHAT_ID, "text": msg, "parse_mode": "Markdown"},
                timeout=10.0
            )
            print(f"Telegram alert sent, status: {res.status_code}", flush=True)
    except Exception as e:
        print(f"Failed to send Telegram alert: {e}", flush=True)

@router.post("/login", response_model=Token)
@limiter.limit("5/minute")
async def login(creds: LoginRequest, request: Request):
    client_ip = request.client.host
    # Verify username/password
    if creds.username != config.ADMIN_USER or not auth.verify_password(creds.password, auth.load_password_hash()):
        asyncio.create_task(send_failed_login_alert(client_ip, creds.username, "Incorrect username or password"))
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Incorrect username or password",
            headers={"WWW-Authenticate": "Bearer"},
        )
    
    # Check 2FA
    totp_secret = auth.load_totp_secret()
    if totp_secret:
        if not creds.totp_code:
            asyncio.create_task(send_failed_login_alert(client_ip, creds.username, "Missing 2FA Code"))
            raise HTTPException(
                status_code=status.HTTP_403_FORBIDDEN,
                detail="2FA code required",
                headers={"WWW-Authenticate": "Bearer"},
            )
        if not auth.verify_totp(creds.totp_code, totp_secret):
            asyncio.create_task(send_failed_login_alert(client_ip, creds.username, "Invalid 2FA Code"))
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail="Invalid 2FA code",
                headers={"WWW-Authenticate": "Bearer"},
            )

    access_token = auth.create_access_token(
        data={"sub": creds.username},
        expires_delta=timedelta(minutes=config.ACCESS_TOKEN_EXPIRE_MINUTES),
    )
    refresh_token = auth.create_refresh_token(data={"sub": creds.username})
    return {
        "access_token": access_token,
        "refresh_token": refresh_token,
        "token_type": "bearer",
        "user": creds.username,
        "has_2fa": bool(totp_secret)
    }

@router.get("/me")
async def read_users_me(current_user: str = Depends(auth.get_current_user)):
    totp_secret = auth.load_totp_secret()
    return {"username": current_user, "has_2fa": bool(totp_secret)}

@router.post("/refresh")
async def refresh_token(req: RefreshRequest):
    try:
        payload = jwt.decode(req.refresh_token, config.SECRET_KEY, algorithms=[config.ALGORITHM])
        if payload.get("type") != "refresh":
            raise HTTPException(status_code=401, detail="Invalid token type")
        username = payload.get("sub")
        if username != config.ADMIN_USER:
            raise HTTPException(status_code=401, detail="Invalid user")
    except jwt.PyJWTError:
        raise HTTPException(status_code=401, detail="Invalid refresh token")
    access_token = auth.create_access_token(
        data={"sub": username},
        expires_delta=timedelta(minutes=config.ACCESS_TOKEN_EXPIRE_MINUTES),
    )
    refresh_token = auth.create_refresh_token(data={"sub": username})
    return {"access_token": access_token, "refresh_token": refresh_token}

@router.post("/setup-2fa/generate", response_model=Setup2FAResponse)
async def generate_2fa(current_user: str = Depends(auth.get_current_user)):
    secret = pyotp.random_base32()
    uri = pyotp.totp.TOTP(secret).provisioning_uri(name=current_user, issuer_name="NAS Dashboard")
    return {"secret": secret, "uri": uri}

@router.post("/setup-2fa/verify")
async def verify_2fa_setup(request: Verify2FARequest, current_user: str = Depends(auth.get_current_user)):
    totp = pyotp.TOTP(request.secret)
    if not totp.verify(request.code):
        raise HTTPException(status_code=400, detail="Invalid code")
    
    # Save secret
    auth.save_totp_secret(request.secret)
    return {"message": "2FA enabled successfully"}
    
@router.post("/setup-2fa/disable")
async def disable_2fa(current_user: str = Depends(auth.get_current_user)):
    auth.save_totp_secret("")
    return {"message": "2FA disabled"}

class ChangePasswordRequest(BaseModel):
    current_password: str
    new_password: str

@router.post("/change-password")
@limiter.limit("5/minute")
async def change_password(req: ChangePasswordRequest, request: Request, current_user: str = Depends(auth.get_current_user)):
    if not auth.verify_password(req.current_password, auth.load_password_hash()):
        raise HTTPException(status_code=400, detail="Current password is incorrect")
    if len(req.new_password) < 8:
        raise HTTPException(status_code=400, detail="Password must be at least 8 characters")
    auth.save_password_hash(auth.get_password_hash(req.new_password))
    return {"message": "Password changed successfully"}

# WebAuthn — in-memory challenge store with TTL
_challenges = {}

def _store_challenge(challenge: bytes):
    _challenges["current"] = (challenge, time.time())

def _get_challenge() -> bytes:
    entry = _challenges.pop("current", None)
    if not entry or time.time() - entry[1] > 300:
        raise HTTPException(status_code=400, detail="Challenge expired")
    return entry[0]

@router.post("/passkey/register/options")
async def passkey_register_options(current_user: str = Depends(auth.get_current_user)):
    existing = auth.load_passkey_credentials()
    exclude = [PublicKeyCredentialDescriptor(id=base64url_to_bytes(c["credential_id"])) for c in existing]
    options = generate_registration_options(
        rp_id=config.WEBAUTHN_RP_ID,
        rp_name="NAS Dashboard",
        user_id=config.ADMIN_USER.encode(),
        user_name=config.ADMIN_USER,
        user_display_name=config.ADMIN_USER,
        exclude_credentials=exclude,
    )
    _store_challenge(options.challenge)
    return JSONResponse(content=json.loads(options_to_json(options)))

@router.post("/passkey/register/verify")
async def passkey_register_verify(request: Request, current_user: str = Depends(auth.get_current_user)):
    body = await request.json()
    challenge = _get_challenge()
    try:
        verification = verify_registration_response(
            credential=body,
            expected_challenge=challenge,
            expected_rp_id=config.WEBAUTHN_RP_ID,
            expected_origin=config.WEBAUTHN_ORIGIN,
        )
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))
    auth.save_passkey_credential({
        "credential_id": bytes_to_base64url(verification.credential_id),
        "public_key": bytes_to_base64url(verification.credential_public_key),
        "sign_count": verification.sign_count,
        "name": body.get("name", "Passkey"),
        "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
    })
    return {"ok": True}

@router.post("/passkey/login/options")
async def passkey_login_options():
    creds = auth.load_passkey_credentials()
    if not creds:
        raise HTTPException(status_code=404, detail="No passkeys registered")
    allow = [PublicKeyCredentialDescriptor(id=base64url_to_bytes(c["credential_id"])) for c in creds]
    options = generate_authentication_options(
        rp_id=config.WEBAUTHN_RP_ID,
        allow_credentials=allow,
        user_verification=UserVerificationRequirement.PREFERRED,
    )
    _store_challenge(options.challenge)
    return JSONResponse(content=json.loads(options_to_json(options)))

@router.post("/passkey/login/verify")
async def passkey_login_verify(request: Request):
    body = await request.json()
    challenge = _get_challenge()
    creds = auth.load_passkey_credentials()
    cred_id_b64 = body.get("id", "")
    matched = next((c for c in creds if c["credential_id"] == cred_id_b64), None)
    if not matched:
        raise HTTPException(status_code=400, detail="Unknown credential")
    try:
        verification = verify_authentication_response(
            credential=body,
            expected_challenge=challenge,
            expected_rp_id=config.WEBAUTHN_RP_ID,
            expected_origin=config.WEBAUTHN_ORIGIN,
            credential_public_key=base64url_to_bytes(matched["public_key"]),
            credential_current_sign_count=matched["sign_count"],
        )
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))
    matched["sign_count"] = verification.new_sign_count
    data = auth._load_auth_data()
    data["passkey_credentials"] = creds
    auth._save_auth_data(data)
    username = config.ADMIN_USER
    access_token = auth.create_access_token(
        data={"sub": username},
        expires_delta=timedelta(minutes=config.ACCESS_TOKEN_EXPIRE_MINUTES),
    )
    refresh_token = auth.create_refresh_token(data={"sub": username})
    return {"access_token": access_token, "refresh_token": refresh_token, "token_type": "bearer", "user": username}

@router.get("/passkey/list")
async def passkey_list(current_user: str = Depends(auth.get_current_user)):
    creds = auth.load_passkey_credentials()
    return {"passkeys": [{"id": c["credential_id"], "name": c.get("name", "Passkey"), "created_at": c.get("created_at", "")} for c in creds]}

@router.post("/passkey/delete")
async def passkey_delete(request: Request, current_user: str = Depends(auth.get_current_user)):
    body = await request.json()
    cred_id = body.get("id")
    data = auth._load_auth_data()
    creds = data.get("passkey_credentials", [])
    data["passkey_credentials"] = [c for c in creds if c["credential_id"] != cred_id]
    auth._save_auth_data(data)
    return {"ok": True}

@router.post("/passkey/clear")
async def passkey_clear(current_user: str = Depends(auth.get_current_user)):
    auth.clear_passkey_credentials()
    return {"ok": True}