Gan, Jimmy
136 lines
5.2 KiB
Python
136 lines
5.2 KiB
Python
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"}
|