feat: 增加可信应用免登录接入
This commit is contained in:
213
ai_knowledge_base_v2/apps/backend/tests/test_sso_service.py
Normal file
213
ai_knowledge_base_v2/apps/backend/tests/test_sso_service.py
Normal file
@@ -0,0 +1,213 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import hmac
|
||||
import json
|
||||
import secrets
|
||||
import time
|
||||
|
||||
from fastapi import HTTPException
|
||||
import pytest
|
||||
from sqlalchemy import create_engine, func, select
|
||||
from sqlalchemy.orm import Session
|
||||
from sqlalchemy.pool import StaticPool
|
||||
|
||||
from app.models import Base
|
||||
from app.core.config import get_settings
|
||||
from app.models.ai_config import SystemConfig
|
||||
from app.models.sso import SsoClient, SsoLoginAudit, UserExternalIdentity
|
||||
from app.models.user import User
|
||||
from app.schemas.sso import SsoTicketRequest
|
||||
from app.services.secret_service import SecretService
|
||||
from app.services.sso_service import SsoService
|
||||
|
||||
|
||||
class FakeRedis:
|
||||
def __init__(self):
|
||||
self.values: dict[str, str] = {}
|
||||
|
||||
def set(self, key: str, value: str, *, ex: int, nx: bool = False):
|
||||
if nx and key in self.values:
|
||||
return False
|
||||
self.values[key] = value
|
||||
return True
|
||||
|
||||
def eval(self, _script: str, _key_count: int, key: str):
|
||||
return self.values.pop(key, None)
|
||||
|
||||
|
||||
def _db() -> Session:
|
||||
engine = create_engine(
|
||||
"sqlite:///:memory:",
|
||||
connect_args={"check_same_thread": False},
|
||||
poolclass=StaticPool,
|
||||
)
|
||||
Base.metadata.create_all(engine)
|
||||
return Session(engine)
|
||||
|
||||
|
||||
def _signature(secret: str, timestamp: str, nonce: str, body: bytes) -> str:
|
||||
body_hash = hashlib.sha256(body).hexdigest()
|
||||
canonical = f"{timestamp}\n{nonce}\n{body_hash}".encode()
|
||||
return hmac.new(secret.encode(), canonical, hashlib.sha256).hexdigest()
|
||||
|
||||
|
||||
def _seed(db: Session) -> tuple[User, SsoClient, str]:
|
||||
secret = "integration-secret"
|
||||
user = User(
|
||||
id=1,
|
||||
phone="13800138000",
|
||||
name="测试学员",
|
||||
status=1,
|
||||
daily_chat_limit=100,
|
||||
daily_chat_used=0,
|
||||
is_deleted=0,
|
||||
)
|
||||
client = SsoClient(
|
||||
id=1,
|
||||
app_id="student-app",
|
||||
name="学员应用",
|
||||
client_secret=SecretService.encrypt(secret),
|
||||
redirect_uris='["https://student.example.com/home"]',
|
||||
status=1,
|
||||
)
|
||||
db.add_all(
|
||||
[
|
||||
user,
|
||||
client,
|
||||
SystemConfig(config_key="sso_user_client_url", config_value="https://qa.example.com"),
|
||||
]
|
||||
)
|
||||
db.commit()
|
||||
return user, client, secret
|
||||
|
||||
|
||||
def test_sso_ticket_binds_user_and_exchanges_only_once(monkeypatch: pytest.MonkeyPatch):
|
||||
redis = FakeRedis()
|
||||
monkeypatch.setattr("app.services.sso_service.get_sync_redis_client", lambda: redis)
|
||||
monkeypatch.setattr(get_settings(), "jwt_secret_key", "test-sso-jwt-secret-key-32-bytes-long")
|
||||
with _db() as db:
|
||||
user, client, secret = _seed(db)
|
||||
raw_body = json.dumps(
|
||||
{
|
||||
"externalUserId": "external-1001",
|
||||
"verifiedPhone": user.phone,
|
||||
"displayName": "外部昵称",
|
||||
"returnUrl": "https://student.example.com/home",
|
||||
},
|
||||
ensure_ascii=False,
|
||||
separators=(",", ":"),
|
||||
).encode()
|
||||
payload = SsoTicketRequest.model_validate_json(raw_body)
|
||||
timestamp = str(int(time.time()))
|
||||
nonce = secrets.token_urlsafe(24)
|
||||
|
||||
result = SsoService.issue_ticket(
|
||||
db,
|
||||
app_id=client.app_id,
|
||||
timestamp=timestamp,
|
||||
nonce=nonce,
|
||||
signature=_signature(secret, timestamp, nonce, raw_body),
|
||||
raw_body=raw_body,
|
||||
payload=payload,
|
||||
ip="127.0.0.1",
|
||||
)
|
||||
|
||||
assert result["entryUrl"].startswith("https://qa.example.com/?sso_code=")
|
||||
assert db.scalar(select(func.count(UserExternalIdentity.id))) == 1
|
||||
|
||||
login = SsoService.exchange(db, code=result["code"], ip="127.0.0.1")
|
||||
assert login["user"].id == user.id
|
||||
assert login["returnUrl"] == "https://student.example.com/home"
|
||||
assert login["token"]
|
||||
|
||||
with pytest.raises(HTTPException) as reused:
|
||||
SsoService.exchange(db, code=result["code"], ip="127.0.0.1")
|
||||
assert reused.value.status_code == 401
|
||||
assert db.scalar(select(func.count(SsoLoginAudit.id))) == 3
|
||||
|
||||
|
||||
def test_sso_rejects_replayed_nonce_and_unlisted_return_url(monkeypatch: pytest.MonkeyPatch):
|
||||
redis = FakeRedis()
|
||||
monkeypatch.setattr("app.services.sso_service.get_sync_redis_client", lambda: redis)
|
||||
with _db() as db:
|
||||
user, client, secret = _seed(db)
|
||||
timestamp = str(int(time.time()))
|
||||
nonce = "nonce-used-only-once-123"
|
||||
raw_body = json.dumps(
|
||||
{"externalUserId": "external-1001", "verifiedPhone": user.phone},
|
||||
separators=(",", ":"),
|
||||
).encode()
|
||||
payload = SsoTicketRequest.model_validate_json(raw_body)
|
||||
signature = _signature(secret, timestamp, nonce, raw_body)
|
||||
SsoService.issue_ticket(
|
||||
db,
|
||||
app_id=client.app_id,
|
||||
timestamp=timestamp,
|
||||
nonce=nonce,
|
||||
signature=signature,
|
||||
raw_body=raw_body,
|
||||
payload=payload,
|
||||
ip="127.0.0.1",
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as replayed:
|
||||
SsoService.issue_ticket(
|
||||
db,
|
||||
app_id=client.app_id,
|
||||
timestamp=timestamp,
|
||||
nonce=nonce,
|
||||
signature=signature,
|
||||
raw_body=raw_body,
|
||||
payload=payload,
|
||||
ip="127.0.0.1",
|
||||
)
|
||||
assert replayed.value.status_code == 409
|
||||
|
||||
another_nonce = "another-unique-nonce-456"
|
||||
invalid_body = json.dumps(
|
||||
{
|
||||
"externalUserId": "external-1001",
|
||||
"verifiedPhone": user.phone,
|
||||
"returnUrl": "https://evil.example.com",
|
||||
},
|
||||
separators=(",", ":"),
|
||||
).encode()
|
||||
with pytest.raises(HTTPException) as invalid_return:
|
||||
SsoService.issue_ticket(
|
||||
db,
|
||||
app_id=client.app_id,
|
||||
timestamp=timestamp,
|
||||
nonce=another_nonce,
|
||||
signature=_signature(secret, timestamp, another_nonce, invalid_body),
|
||||
raw_body=invalid_body,
|
||||
payload=SsoTicketRequest.model_validate_json(invalid_body),
|
||||
ip="127.0.0.1",
|
||||
)
|
||||
assert invalid_return.value.status_code == 400
|
||||
|
||||
|
||||
def test_sso_first_login_requires_verified_existing_phone(monkeypatch: pytest.MonkeyPatch):
|
||||
redis = FakeRedis()
|
||||
monkeypatch.setattr("app.services.sso_service.get_sync_redis_client", lambda: redis)
|
||||
with _db() as db:
|
||||
_user, client, secret = _seed(db)
|
||||
raw_body = b'{"externalUserId":"external-without-phone"}'
|
||||
payload = SsoTicketRequest.model_validate_json(raw_body)
|
||||
timestamp = str(int(time.time()))
|
||||
nonce = "missing-phone-nonce-123"
|
||||
with pytest.raises(HTTPException) as missing_phone:
|
||||
SsoService.issue_ticket(
|
||||
db,
|
||||
app_id=client.app_id,
|
||||
timestamp=timestamp,
|
||||
nonce=nonce,
|
||||
signature=_signature(secret, timestamp, nonce, raw_body),
|
||||
raw_body=raw_body,
|
||||
payload=payload,
|
||||
ip="127.0.0.1",
|
||||
)
|
||||
assert missing_phone.value.status_code == 403
|
||||
audit = db.scalar(select(SsoLoginAudit).order_by(SsoLoginAudit.id.desc()))
|
||||
assert audit is not None
|
||||
assert audit.status == "FAILED"
|
||||
Reference in New Issue
Block a user