from datetime import timedelta import asyncio import httpx from fastapi import APIRouter, Depends, HTTPException, status, Request from pydantic import BaseModel from slowapi import Limiter from slowapi.util import get_remote_address import config import auth import pyotp 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 token_type: str user: str has_2fa: bool 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), ) return { "access_token": access_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("/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"}