102 lines
2.7 KiB
Python
102 lines
2.7 KiB
Python
"""2FA TOTP service for VoIdea."""
|
|
|
|
import base64
|
|
import io
|
|
import json
|
|
from typing import Any
|
|
|
|
import pyotp
|
|
import qrcode
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.models.agent import AgentConfig
|
|
from app.models.user import User
|
|
|
|
|
|
def generate_totp_secret() -> str:
|
|
return pyotp.random_base32()
|
|
|
|
|
|
def get_totp_uri(secret: str, email: str, issuer: str = "VoIdeaAI") -> str:
|
|
return pyotp.totp.TOTP(secret).provisioning_uri(name=email, issuer_name=issuer)
|
|
|
|
|
|
def generate_qr_base64(uri: str) -> str:
|
|
qr = qrcode.make(uri)
|
|
buf = io.BytesIO()
|
|
qr.save(buf, format="PNG")
|
|
return base64.b64encode(buf.getvalue()).decode()
|
|
|
|
|
|
def verify_totp(secret: str, token: str) -> bool:
|
|
totp = pyotp.TOTP(secret)
|
|
return totp.verify(token)
|
|
|
|
|
|
def get_user_secret(user: User) -> str | None:
|
|
tuning = user.pipeline_tuning or {}
|
|
return tuning.get("totp_secret")
|
|
|
|
|
|
def set_user_secret(user: User, secret: str) -> None:
|
|
tuning = user.pipeline_tuning or {}
|
|
tuning["totp_secret"] = secret
|
|
user.pipeline_tuning = tuning
|
|
|
|
|
|
def is_2fa_enabled(user: User) -> bool:
|
|
tuning = user.pipeline_tuning or {}
|
|
return bool(tuning.get("totp_enabled", False))
|
|
|
|
|
|
def set_2fa_enabled(user: User, enabled: bool) -> None:
|
|
tuning = user.pipeline_tuning or {}
|
|
tuning["totp_enabled"] = enabled
|
|
user.pipeline_tuning = tuning
|
|
|
|
|
|
# ── Global 2FA toggle (stored in AgentConfig for "system") ──
|
|
|
|
_GLOBAL_2FA_KEY = "2fa_globally_enabled"
|
|
|
|
|
|
async def is_2fa_globally_enabled(db: AsyncSession) -> bool:
|
|
result = await db.execute(
|
|
select(AgentConfig).where(AgentConfig.agent_name == "system")
|
|
)
|
|
cfg = result.scalar_one_or_none()
|
|
if not cfg or not cfg.config:
|
|
return True
|
|
try:
|
|
data = json.loads(cfg.config)
|
|
return bool(data.get(_GLOBAL_2FA_KEY, True))
|
|
except (json.JSONDecodeError, TypeError):
|
|
return True
|
|
|
|
|
|
async def set_2fa_globally_enabled(db: AsyncSession, enabled: bool) -> None:
|
|
result = await db.execute(
|
|
select(AgentConfig).where(AgentConfig.agent_name == "system")
|
|
)
|
|
cfg = result.scalar_one_or_none()
|
|
if not cfg:
|
|
cfg = AgentConfig(
|
|
agent_name="system",
|
|
description="System-wide configuration",
|
|
is_enabled=True,
|
|
version="1.0.0",
|
|
config=json.dumps({_GLOBAL_2FA_KEY: enabled}, ensure_ascii=False),
|
|
)
|
|
db.add(cfg)
|
|
else:
|
|
current = {}
|
|
if cfg.config:
|
|
try:
|
|
current = json.loads(cfg.config)
|
|
except (json.JSONDecodeError, TypeError):
|
|
current = {}
|
|
current[_GLOBAL_2FA_KEY] = enabled
|
|
cfg.config = json.dumps(current, ensure_ascii=False)
|
|
await db.commit()
|