Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 50 additions & 0 deletions astrbot/core/desktop_runtime.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,60 @@
import ipaddress
import os
import secrets

DESKTOP_MANAGED_RESTART_MESSAGE = (
"AstrBot Desktop manages this backend process. Please restart or update from "
"the desktop app instead of the core WebUI."
)

DESKTOP_SESSION_SECRET_ENV = "ASTRBOT_DESKTOP_SESSION_SECRET"
DESKTOP_SESSION_SECRET_MIN_LENGTH = 32


def is_desktop_managed_backend() -> bool:
return os.environ.get("ASTRBOT_DESKTOP_MANAGED") == "1"


def get_desktop_session_secret() -> str | None:
"""Return the in-memory secret when desktop session auth is enabled."""
if not is_desktop_managed_backend():
return None

secret = os.environ.get(DESKTOP_SESSION_SECRET_ENV, "").strip()
if len(secret) < DESKTOP_SESSION_SECRET_MIN_LENGTH:
return None
return secret


def is_desktop_session_auth_enabled() -> bool:
return get_desktop_session_secret() is not None


def is_loopback_client_host(host: str | None) -> bool:
if not host:
return False
try:
address = ipaddress.ip_address(host)
except ValueError:
return False

if address.is_loopback:
return True
if isinstance(address, ipaddress.IPv6Address) and address.ipv4_mapped:
return address.ipv4_mapped.is_loopback
return False


def verify_desktop_session_secret(
provided_secret: str | None,
client_host: str | None,
) -> bool:
configured_secret = get_desktop_session_secret()
if configured_secret is None or not is_loopback_client_host(client_host):
return False
if not isinstance(provided_secret, str):
return False
return secrets.compare_digest(
provided_secret.encode("utf-8"),
configured_secret.encode("utf-8"),
)
14 changes: 14 additions & 0 deletions astrbot/dashboard/api/auth.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,8 @@
AuthServiceResult,
)

DESKTOP_SESSION_HEADER = "X-AstrBot-Desktop-Session"

router = APIRouter(tags=["Auth"])
legacy_router = APIRouter(
prefix="/api/auth",
Expand Down Expand Up @@ -389,6 +391,18 @@ async def login(
return await _login(request, payload, service)


@router.post("/auth/desktop-session", include_in_schema=False)
async def desktop_session(
request: Request,
service: AuthService = Depends(get_auth_service),
):
result = await service.desktop_session(
request.headers.get(DESKTOP_SESSION_HEADER),
request.client.host if request.client else None,
)
return _auth_service_response(request, result)


@legacy_router.post("/login")
async def dashboard_login(
request: Request,
Expand Down
1 change: 1 addition & 0 deletions astrbot/dashboard/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@
"/api/v1/auth/totp/setup",
"/api/auth/login",
"/api/v1/auth/login",
"/api/v1/auth/desktop-session",
}
)

Expand Down
42 changes: 41 additions & 1 deletion astrbot/dashboard/services/auth_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,11 @@
from astrbot.core import DEMO_MODE
from astrbot.core.config.astrbot_config import AstrBotConfig
from astrbot.core.db import BaseDatabase
from astrbot.core.desktop_runtime import (
is_desktop_session_auth_enabled,
is_loopback_client_host,
verify_desktop_session_secret,
)
from astrbot.core.utils.auth_password import (
is_default_dashboard_password,
is_md5_dashboard_password,
Expand Down Expand Up @@ -122,6 +127,14 @@ def __init__(
self.demo_mode = demo_mode

async def setup_status(self) -> AuthServiceResult:
if is_desktop_session_auth_enabled():
return AuthServiceResult(
data={
"setup_required": False,
"skip_default_password_auth": False,
"password_upgrade_required": False,
}
)
return AuthServiceResult(
data={
"setup_required": await self.is_setup_required(),
Expand All @@ -133,6 +146,31 @@ async def setup_status(self) -> AuthServiceResult:
}
)

async def desktop_session(
self,
provided_secret: str | None,
client_host: str | None,
) -> AuthServiceResult:
if not is_desktop_session_auth_enabled() or not is_loopback_client_host(
client_host
):
return self.error("Not found", status_code=404)
if not verify_desktop_session_secret(provided_secret, client_host):
return self.error("Invalid desktop session", status_code=401)

username = self.config["dashboard"]["username"]
token = self.generate_jwt(username, auth_source="desktop")
return AuthServiceResult(
data={
"token": token,
"username": username,
"change_pwd_hint": False,
"md5_pwd_hint": False,
"password_upgrade_required": False,
},
jwt_token=token,
)

async def totp_setup(self, post_data: object) -> AuthServiceResult:
if isinstance(post_data, dict) and post_data.get("secret"):
secret = post_data["secret"]
Expand Down Expand Up @@ -406,12 +444,14 @@ async def edit_account(self, post_data: object) -> AuthServiceResult:

return AuthServiceResult(message="Updated account successfully")

def generate_jwt(self, username: str):
def generate_jwt(self, username: str, *, auth_source: str = "password"):
payload = {
"username": username,
"exp": datetime.datetime.now(datetime.timezone.utc)
+ datetime.timedelta(days=7),
}
if auth_source != "password":
payload["auth_source"] = auth_source
jwt_token = self.config["dashboard"].get("jwt_secret", None)
if not jwt_token:
raise ValueError("JWT secret is not set in the cmd_config.")
Expand Down
11 changes: 11 additions & 0 deletions astrbot/dashboard/services/stat_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
from astrbot.core.desktop_runtime import (
DESKTOP_MANAGED_RESTART_MESSAGE,
is_desktop_managed_backend,
is_desktop_session_auth_enabled,
)
from astrbot.core.utils.astrbot_path import get_astrbot_path
from astrbot.core.utils.auth_password import (
Expand Down Expand Up @@ -75,6 +76,8 @@ def get_running_time_components(total_seconds: int):
return {"hours": hours, "minutes": minutes, "seconds": seconds}

async def is_default_cred(self):
if is_desktop_session_auth_enabled():
return False
password_change_required = await is_password_change_required(
self.db_helper,
self.config,
Expand All @@ -96,6 +99,14 @@ async def is_default_cred(self):
) and not DEMO_MODE

async def get_version(self) -> dict:
if is_desktop_session_auth_enabled():
return {
"version": VERSION,
"dashboard_version": await get_dashboard_version(),
"change_pwd_hint": False,
"md5_pwd_hint": False,
"password_upgrade_required": False,
}
storage_upgraded = await is_password_storage_upgraded(
self.db_helper,
self.config,
Expand Down
115 changes: 115 additions & 0 deletions tests/test_dashboard.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
from types import SimpleNamespace
from urllib.parse import parse_qs, urlsplit, urlunsplit

import jwt
import pyotp
import pytest
import pytest_asyncio
Expand Down Expand Up @@ -435,6 +436,120 @@ async def test_auth_login(
assert "Secure" not in jwt_cookie_header


@pytest.mark.asyncio
async def test_desktop_session_issues_jwt_without_password(
app: FastAPIAppAdapter,
core_lifecycle_td: AstrBotCoreLifecycle,
monkeypatch: pytest.MonkeyPatch,
):
secret = "desktop-session-secret-" * 2
monkeypatch.setenv("ASTRBOT_DESKTOP_MANAGED", "1")
monkeypatch.setenv("ASTRBOT_DESKTOP_SESSION_SECRET", secret)
app._dashboard_server._rate_limiter_registry.clear()

response = await app.test_client().post(
"/api/v1/auth/desktop-session",
headers={"X-AstrBot-Desktop-Session": secret},
json={},
)
data = await response.get_json()

assert response.status_code == 200
assert data["status"] == "ok"
assert data["data"]["username"] == core_lifecycle_td.astrbot_config[
"dashboard"
]["username"]
token = data["data"]["token"]
payload = jwt.decode(
token,
core_lifecycle_td.astrbot_config["dashboard"]["jwt_secret"],
algorithms=["HS256"],
)
assert payload["auth_source"] == "desktop"


@pytest.mark.asyncio
async def test_desktop_session_rejects_wrong_secret(
app: FastAPIAppAdapter,
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setenv("ASTRBOT_DESKTOP_MANAGED", "1")
monkeypatch.setenv("ASTRBOT_DESKTOP_SESSION_SECRET", "a" * 64)
app._dashboard_server._rate_limiter_registry.clear()

response = await app.test_client().post(
"/api/v1/auth/desktop-session",
headers={"X-AstrBot-Desktop-Session": "b" * 64},
json={},
)
data = await response.get_json()

assert response.status_code == 401
assert data["status"] == "error"
assert "token" not in (data.get("data") or {})


@pytest.mark.asyncio
async def test_desktop_session_endpoint_is_hidden_when_not_managed(
app: FastAPIAppAdapter,
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.delenv("ASTRBOT_DESKTOP_MANAGED", raising=False)
monkeypatch.delenv("ASTRBOT_DESKTOP_SESSION_SECRET", raising=False)
app._dashboard_server._rate_limiter_registry.clear()

response = await app.test_client().post(
"/api/v1/auth/desktop-session",
headers={"X-AstrBot-Desktop-Session": "a" * 64},
json={},
)
data = await response.get_json()

assert response.status_code == 404
assert data["status"] == "error"


@pytest.mark.asyncio
async def test_desktop_session_suppresses_password_setup_and_warnings(
app: FastAPIAppAdapter,
core_lifecycle_td: AstrBotCoreLifecycle,
monkeypatch: pytest.MonkeyPatch,
):
secret = "desktop-session-secret-" * 2
monkeypatch.setenv("ASTRBOT_DESKTOP_MANAGED", "1")
monkeypatch.setenv("ASTRBOT_DESKTOP_SESSION_SECRET", secret)
app._dashboard_server._rate_limiter_registry.clear()
await _set_dashboard_password_change_required(core_lifecycle_td, True)

try:
client = app.test_client()
setup_response = await client.get("/api/v1/auth/setup-status")
setup_data = await setup_response.get_json()
assert setup_data["data"] == {
"setup_required": False,
"skip_default_password_auth": False,
"password_upgrade_required": False,
}

session_response = await client.post(
"/api/v1/auth/desktop-session",
headers={"X-AstrBot-Desktop-Session": secret},
json={},
)
session_data = await session_response.get_json()
token = session_data["data"]["token"]
version_response = await client.get(
"/api/v1/stats/version",
headers={"Authorization": f"Bearer {token}"},
)
version_data = await version_response.get_json()
assert version_data["data"]["change_pwd_hint"] is False
assert version_data["data"]["md5_pwd_hint"] is False
assert version_data["data"]["password_upgrade_required"] is False
finally:
await _set_dashboard_password_change_required(core_lifecycle_td, False)


@pytest.mark.asyncio
async def test_auth_login_secure_cookie_override(
app: FastAPIAppAdapter,
Expand Down
42 changes: 42 additions & 0 deletions tests/test_desktop_runtime.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,42 @@
from astrbot.core.desktop_runtime import (
DESKTOP_SESSION_SECRET_ENV,
get_desktop_session_secret,
is_desktop_session_auth_enabled,
is_loopback_client_host,
verify_desktop_session_secret,
)


def test_desktop_session_auth_requires_managed_backend(monkeypatch):
monkeypatch.delenv("ASTRBOT_DESKTOP_MANAGED", raising=False)
monkeypatch.setenv(DESKTOP_SESSION_SECRET_ENV, "a" * 64)

assert get_desktop_session_secret() is None
assert is_desktop_session_auth_enabled() is False


def test_desktop_session_auth_rejects_short_secret(monkeypatch):
monkeypatch.setenv("ASTRBOT_DESKTOP_MANAGED", "1")
monkeypatch.setenv(DESKTOP_SESSION_SECRET_ENV, "too-short")

assert get_desktop_session_secret() is None
assert is_desktop_session_auth_enabled() is False


def test_desktop_session_secret_only_matches_on_loopback(monkeypatch):
secret = "desktop-session-secret-" * 2
monkeypatch.setenv("ASTRBOT_DESKTOP_MANAGED", "1")
monkeypatch.setenv(DESKTOP_SESSION_SECRET_ENV, secret)

assert verify_desktop_session_secret(secret, "127.0.0.1") is True
assert verify_desktop_session_secret(secret, "::1") is True
assert verify_desktop_session_secret(secret, "::ffff:127.0.0.1") is True
assert verify_desktop_session_secret("wrong" * 10, "127.0.0.1") is False
assert verify_desktop_session_secret(secret, "192.168.1.10") is False


def test_loopback_client_host_rejects_names_and_unspecified_addresses():
assert is_loopback_client_host("localhost") is False
assert is_loopback_client_host("0.0.0.0") is False
assert is_loopback_client_host("::") is False
assert is_loopback_client_host(None) is False
Loading