MCP Server 的身份认证方案

# mcp_mfa_service.py
from typing import List, Dict, Optional
from enum import Enum
from datetime import datetime, timedelta
import random
import string
import logging

class MFAType(Enum):
    SMS = "sms"
    EMAIL = "email"
    TOTP = "totp"
    PUSH = "push"
    BIOMETRIC = "biometric"

class MFAChallenge:
    def __init__(self, challenge_id: str, user_id: str, mfa_type: MFAType, 
                 challenge_data: Dict, expires_at: datetime, attempts: int = 0):
        self.challenge_id = challenge_id
        self.user_id = user_id
        self.mfa_type = mfa_type
        self.challenge_data = challenge_data
        self.expires_at = expires_at
        self.attempts = attempts
    
    def to_dict(self) -> Dict:
        return {
            "challenge_id": self.challenge_id,
            "user_id": self.user_id,
            "mfa_type": self.mfa_type.value,
            "challenge_data": self.challenge_data,
            "expires_at": self.expires_at.isoformat(),
            "attempts": self.attempts
        }
    
    def is_expired(self) -> bool:
        return datetime.now() > self.expires_at
    
    def increment_attempts(self):
        self.attempts += 1

class MFAService:
    def __init__(self):
        self.challenges = {}
        self.logger = logging.getLogger("mcp_mfa")
    
    def generate_challenge(self, user_id: str, mfa_type: MFAType, 
                         challenge_data: Optional[Dict] = None) -> MFAChallenge:
        """生成多因素认证挑战"""
        challenge_id = self._generate_unique_id()
        expires_at = datetime.now() + timedelta(minutes=5)
        
        if not challenge_data:
            challenge_data = self._generate_challenge_data(mfa_type)
        
        challenge = MFAChallenge(
            challenge_id=challenge_id,
            user_id=user_id,
            mfa_type=mfa_type,
            challenge_data=challenge_data,
            expires_at=expires_at
        )
        
        self.challenges[challenge_id] = challenge
        self.logger.info(f"Generated MFA challenge: {challenge_id} for user: {user_id}")
        return challenge
    
    def verify_challenge(self, challenge_id: str, response: str) -> bool:
        """验证多因素认证挑战"""
        if challenge_id not in self.challenges:
            self.logger.error(f"Challenge not found: {challenge_id}")
            return False
        
        challenge = self.challenges[challenge_id]
        
        # 检查挑战是否过期
        if challenge.is_expired():
            self.logger.error(f"Challenge expired: {challenge_id}")
            del self.challenges[challenge_id]
            return False
        
        # 检查尝试次数是否超过限制
        if challenge.attempts >= 3:
            self.logger.error(f"Too many attempts for challenge: {challenge_id}")
            del self.challenges[challenge_id]
            return False
        
        # 验证响应
        if self._verify_response(challenge, response):
            self.logger.info(f"Challenge verified successfully: {challenge_id}")
            del self.challenges[challenge_id]
            return True
        else:
            challenge.increment_attempts()
            self.logger.error(f"Challenge verification failed: {challenge_id}")
            return False
    
    def _generate_unique_id(self) -> str:
        """生成唯一 ID"""
        return ''.join(random.choices(string.ascii_letters + string.digits, k=32))
    
    def _generate_challenge_data(self, mfa_type: MFAType) -> Dict:
        """生成挑战数据"""
        if mfa_type == MFAType.SMS or mfa_type == MFAType.EMAIL:
            # 生成 6 位数字验证码
            code = ''.join(random.choices(string.digits, k=6))
            return {"code": code}
        elif mfa_type == MFAType.TOTP:
            # TOTP 不需要预先生成挑战数据
            return {}
        elif mfa_type == MFAType.PUSH:
            # 推送通知挑战数据
            return {"message": "请确认登录请求"}
        elif mfa_type == MFAType.BIOMETRIC:
            # 生物特征挑战数据
            return {"biometric_type": "fingerprint"}
        else:
            return {}
    
    def _verify_response(self, challenge: MFAChallenge, response: str) -> bool:
        """验证响应"""
        if challenge.mfa_type == MFAType.SMS or challenge.mfa_type == MFAType.EMAIL:
            return challenge.challenge_data.get("code") == response
        elif challenge.mfa_type == MFAType.TOTP:
            # 实现 TOTP 验证逻辑
            return self._verify_totp(challenge.user_id, response)
        elif challenge.mfa_type == MFAType.PUSH:
            # 实现推送通知验证逻辑
            return response == "approved"
        elif challenge.mfa_type == MFAType.BIOMETRIC:
            # 实现生物特征验证逻辑
            return self._verify_biometric(challenge.user_id, response)
        else:
            return False
    
    def _verify_totp(self, user_id: str, code: str) -> bool:
        """验证 TOTP 代码"""
        # 实现 TOTP 验证逻辑
        # 这里使用 pyotp 库实现 TOTP 验证
        import pyotp
        
        # 从用户管理服务获取用户的 TOTP 密钥
        totp_secret = self._get_user_totp_secret(user_id)
        if not totp_secret:
            return False
        
        totp = pyotp.TOTP(totp_secret)
        return totp.verify(code)
    
    def _verify_biometric(self, user_id: str, biometric_data: str) -> bool:
        """验证生物特征数据"""
        # 实现生物特征验证逻辑
        # 这里简化处理，实际应调用生物特征验证服务
        return True
    
    def _get_user_totp_secret(self, user_id: str) -> str:
        """获取用户的 TOTP 密钥"""
        # 从用户管理服务获取用户的 TOTP 密钥
        # 这里简化处理，实际应调用用户管理服务
        return "JBSWY3DPEHPK3PXP"
    
    def send_challenge(self, challenge: MFAChallenge):
        """发送多因素认证挑战"""
        if challenge.mfa_type == MFAType.SMS:
            self._send_sms(challenge)
        elif challenge.mfa_type == MFAType.EMAIL:
            self._send_email(challenge)
        elif challenge.mfa_type == MFAType.PUSH:
            self._send_push_notification(challenge)
    
    def _send_sms(self, challenge: MFAChallenge):
        """发送短信验证码"""
        # 实现短信发送逻辑
        code = challenge.challenge_data.get("code")
        self.logger.info(f"Sending SMS with code: {code} to user: {challenge.user_id}")
    
    def _send_email(self, challenge: MFAChallenge):
        """发送邮件验证码"""
        # 实现邮件发送逻辑
        code = challenge.challenge_data.get("code")
        self.logger.info(f"Sending email with code: {code} to user: {challenge.user_id}")
    
    def _send_push_notification(self, challenge: MFAChallenge):
        """发送推送通知"""
        # 实现推送通知发送逻辑
        self.logger.info(f"Sending push notification to user: {challenge.user_id}")

# mcp_oidc_service.py
from typing import Dict, List, Optional
from datetime import datetime, timedelta
import jwt
import logging

class OIDCConfig:
    def __init__(self, issuer: str, client_ids: List[str], 
                 public_key: str, private_key: str, 
                 id_token_expiry: int = 3600, 
                 access_token_expiry: int = 3600):
        self.issuer = issuer
        self.client_ids = client_ids
        self.public_key = public_key
        self.private_key = private_key
        self.id_token_expiry = id_token_expiry
        self.access_token_expiry = access_token_expiry
    
    def to_dict(self) -> Dict:
        return {
            "issuer": self.issuer,
            "client_ids": self.client_ids,
            "id_token_expiry": self.id_token_expiry,
            "access_token_expiry": self.access_token_expiry
        }

class OIDCService:
    def __init__(self, config: OIDCConfig):
        self.config = config
        self.logger = logging.getLogger("mcp_oidc")
    
    def generate_tokens(self, user_id: str, client_id: str, 
                      scopes: List[str] = None) -> Dict:
        """生成 ID Token 和 Access Token"""
        # 检查客户端 ID 是否合法
        if client_id not in self.config.client_ids:
            self.logger.error(f"Invalid client ID: {client_id}")
            return {}
        
        now = datetime.now()
        expires_at = now + timedelta(seconds=self.config.id_token_expiry)
        
        # 生成 ID Token
        id_token = self._generate_id_token(user_id, client_id, scopes, expires_at)
        
        # 生成 Access Token
        access_token = self._generate_access_token(user_id, client_id, scopes, expires_at)
        
        self.logger.info(f"Generated tokens for user: {user_id}, client: {client_id}")
        return {
            "id_token": id_token,
            "access_token": access_token,
            "token_type": "Bearer",
            "expires_in": self.config.access_token_expiry
        }
    
    def _generate_id_token(self, user_id: str, client_id: str, 
                         scopes: List[str], expires_at: datetime) -> str:
        """生成 ID Token"""
        claims = {
            "iss": self.config.issuer,
            "sub": user_id,
            "aud": client_id,
            "iat": datetime.now(),
            "exp": expires_at,
            "auth_time": datetime.now(),
            "nonce": self._generate_nonce(),
            "acr": "urn:mcp:loa:2",  # MCP 认证等级
            "amr": ["pwd", "mfa"]  # 认证方法参考
        }
        
        # 添加自定义声明
        claims["mcp_user_type"] = self._get_user_type(user_id)
        
        # 使用私钥签名 ID Token
        id_token = jwt.encode(
            claims=claims,
            key=self.config.private_key,
            algorithm="RS256"
        )
        
        return id_token
    
    def _generate_access_token(self, user_id: str, client_id: str, 
                             scopes: List[str], expires_at: datetime) -> str:
        """生成 Access Token"""
        claims = {
            "iss": self.config.issuer,
            "sub": user_id,
            "aud": client_id,
            "iat": datetime.now(),
            "exp": expires_at + timedelta(seconds=self.config.access_token_expiry - self.config.id_token_expiry),
            "scope": " ".join(scopes or ["openid", "profile", "email"])
        }
        
        # 使用私钥签名 Access Token
        access_token = jwt.encode(
            claims=claims,
            key=self.config.private_key,
            algorithm="RS256"
        )
        
        return access_token
    
    def verify_id_token(self, id_token: str, client_id: str) -> Dict:
        """验证 ID Token"""
        try:
            # 使用公钥验证 ID Token
            claims = jwt.decode(
                jwt=id_token,
                key=self.config.public_key,
                audience=client_id,
                issuer=self.config.issuer,
                algorithms=["RS256"]
            )
            
            self.logger.info(f"Verified ID Token for client: {client_id}")
            return claims
        except jwt.ExpiredSignatureError:
            self.logger.error("ID Token expired")
            return {}
        except jwt.InvalidIssuerError:
            self.logger.error("Invalid issuer")
            return {}
        except jwt.InvalidAudienceError:
            self.logger.error("Invalid audience")
            return {}
        except jwt.InvalidTokenError as e:
            self.logger.error(f"Invalid ID Token: {e}")
            return {}
    
    def verify_access_token(self, access_token: str, client_id: str) -> Dict:
        """验证 Access Token"""
        try:
            # 使用公钥验证 Access Token
            claims = jwt.decode(
                jwt=access_token,
                key=self.config.public_key,
                audience=client_id,
                issuer=self.config.issuer,
                algorithms=["RS256"]
            )
            
            self.logger.info(f"Verified Access Token for client: {client_id}")
            return claims
        except jwt.ExpiredSignatureError:
            self.logger.error("Access Token expired")
            return {}
        except jwt.InvalidIssuerError:
            self.logger.error("Invalid issuer")
            return {}
        except jwt.InvalidAudienceError:
            self.logger.error("Invalid audience")
            return {}
        except jwt.InvalidTokenError as e:
            self.logger.error(f"Invalid Access Token: {e}")
            return {}
    
    def get_user_info(self, access_token: str) -> Dict:
        """获取用户信息"""
        # 验证 Access Token
        claims = self.verify_access_token(access_token, "userinfo")
        if not claims:
            return {}
        
        user_id = claims["sub"]
        
        # 从用户管理服务获取用户信息
        user_info = self._get_user_info(user_id)
        
        return user_info
    
    def _get_user_info(self, user_id: str) -> Dict:
        """获取用户信息"""
        # 实现从用户管理服务获取用户信息的逻辑
        # 这里简化处理，返回模拟数据
        return {
            "sub": user_id,
            "name": "Test User",
            "given_name": "Test",
            "family_name": "User",
            "email": "test@example.com",
            "email_verified": True,
            "preferred_username": "testuser"
        }
    
    def _generate_nonce(self) -> str:
        """生成 Nonce"""
        import secrets
        return secrets.token_urlsafe(16)
    
    def _get_user_type(self, user_id: str) -> str:
        """获取用户类型"""
        # 实现从用户管理服务获取用户类型的逻辑
        # 这里简化处理，返回默认值
        return "user"
    
    def get_discovery_config(self) -> Dict:
        """获取 OIDC 发现配置"""
        return {
            "issuer": self.config.issuer,
            "authorization_endpoint": f"{self.config.issuer}/authorize",
            "token_endpoint": f"{self.config.issuer}/token",
            "userinfo_endpoint": f"{self.config.issuer}/userinfo",
            "jwks_uri": f"{self.config.issuer}/.well-known/jwks.json",
            "response_types_supported": ["code", "token", "id_token"],
            "subject_types_supported": ["public"],
            "id_token_signing_alg_values_supported": ["RS256"],
            "scopes_supported": ["openid", "profile", "email"],
            "token_endpoint_auth_methods_supported": ["client_secret_basic", "client_secret_post"]
        }
    
    def get_jwks(self) -> Dict:
        """获取 JWKS（JSON Web Key Set）"""
        # 实现从公钥生成 JWKS 的逻辑
        # 这里简化处理，返回模拟数据
        return {
            "keys": [
                {
                    "kty": "RSA",
                    "e": "AQAB",
                    "kid": "test-key",
                    "alg": "RS256",
                    "n": "test-public-key"
                }
            ]
        }

# mcp_auth_protection.py
from typing import Dict, Optional
import logging
import re
from datetime import datetime, timedelta

class AuthProtectionService:
    def __init__(self):
        self.failed_login_attempts = {}
        self.lockout_duration = timedelta(minutes=15)
        self.max_failed_attempts = 5
        self.logger = logging.getLogger("mcp_auth_protection")
    
    def check_failed_attempts(self, username: str) -> bool:
        """检查登录失败尝试次数"""
        if username not in self.failed_login_attempts:
            return True
        
        attempts, last_attempt, locked_until = self.failed_login_attempts[username]
        
        # 检查是否已解锁
        if datetime.now() > locked_until:
            # 重置失败尝试次数
            del self.failed_login_attempts[username]
            return True
        
        # 检查是否超过最大失败尝试次数
        if attempts >= self.max_failed_attempts:
            self.logger.warning(f"Account locked: {username} until {locked_until}")
            return False
        
        return True
    
    def record_failed_attempt(self, username: str, ip_address: str):
        """记录登录失败尝试"""
        now = datetime.now()
        
        if username in self.failed_login_attempts:
            attempts, last_attempt, locked_until = self.failed_login_attempts[username]
            attempts += 1
            
            # 如果超过最大失败尝试次数，锁定账号
            if attempts >= self.max_failed_attempts:
                locked_until = now + self.lockout_duration
                self.logger.warning(f"Account locked due to too many failed attempts: {username} from IP: {ip_address}")
            else:
                locked_until = now + self.lockout_duration
        else:
            attempts = 1
            locked_until = now + self.lockout_duration
        
        self.failed_login_attempts[username] = (attempts, now, locked_until)
        self.logger.info(f"Recorded failed login attempt for: {username} from IP: {ip_address}, attempts: {attempts}")
    
    def record_successful_login(self, username: str, ip_address: str):
        """记录登录成功"""
        # 重置失败尝试次数
        if username in self.failed_login_attempts:
            del self.failed_login_attempts[username]
        
        self.logger.info(f"Recorded successful login for: {username} from IP: {ip_address}")
    
    def validate_password_strength(self, password: str) -> Dict:
        """验证密码强度"""
        errors = []
        
        # 检查密码长度
        if len(password) < 8:
            errors.append("密码长度至少为8个字符")
        
        # 检查是否包含大写字母
        if not re.search(r"[A-Z]", password):
            errors.append("密码必须包含至少一个大写字母")
        
        # 检查是否包含小写字母
        if not re.search(r"[a-z]", password):
            errors.append("密码必须包含至少一个小写字母")
        
        # 检查是否包含数字
        if not re.search(r"[0-9]", password):
            errors.append("密码必须包含至少一个数字")
        
        # 检查是否包含特殊字符
        if not re.search(r"[!@#$%^&*(),.?":{}|<>]", password):
            errors.append("密码必须包含至少一个特殊字符")
        
        return {
            "valid": len(errors) == 0,
            "errors": errors
        }
    
    def detect_brute_force(self, ip_address: str) -> bool:
        """检测暴力破解攻击"""
        # 实现基于 IP 地址的暴力破解检测逻辑
        # 这里简化处理，返回默认值
        return False
    
    def detect_phishing_attempt(self, request: Dict) -> bool:
        """检测钓鱼攻击尝试"""
        # 实现钓鱼攻击检测逻辑
        # 这里简化处理，返回默认值
        return False
    
    def generate_csrf_token(self, session_id: str) -> str:
        """生成 CSRF Token"""
        import secrets
        csrf_token = secrets.token_urlsafe(32)
        
        # 将 CSRF Token 存储到会话中
        self._store_csrf_token(session_id, csrf_token)
        
        return csrf_token
    
    def validate_csrf_token(self, session_id: str, csrf_token: str) -> bool:
        """验证 CSRF Token"""
        stored_token = self._get_csrf_token(session_id)
        return stored_token == csrf_token
    
    def _store_csrf_token(self, session_id: str, csrf_token: str):
        """存储 CSRF Token 到会话中"""
        # 实现 CSRF Token 存储逻辑
        pass
    
    def _get_csrf_token(self, session_id: str) -> Optional[str]:
        """从会话中获取 CSRF Token"""
        # 实现 CSRF Token 获取逻辑
        return None

# mcp_auth_api.py
from fastapi import FastAPI, HTTPException, Depends
from pydantic import BaseModel
from typing import List, Optional
import logging

# 导入之前实现的服务
from mcp_mfa_service import MFAService, MFAType
from mcp_oidc_service import OIDCService, OIDCConfig
from mcp_auth_protection import AuthProtectionService

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("mcp_auth_api")

# 创建 FastAPI 应用
app = FastAPI(title="MCP 认证服务 API", version="1.0.0")

# 初始化服务
mfa_service = MFAService()
oauth_config = OIDCConfig(
    issuer="https://auth.mcp.example.com",
    client_ids=["mcp-client-001", "mcp-client-002"],
    public_key="-----BEGIN PUBLIC KEY-----\nMIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA...\n-----END PUBLIC KEY-----",
    private_key="-----BEGIN PRIVATE KEY-----\nMIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggSkAgEAAoIBAQC...\n-----END PRIVATE KEY-----"
)
oidc_service = OIDCService(config=oauth_config)
auth_protection = AuthProtectionService()

# 请求模型
class LoginRequest(BaseModel):
    username: str
    password: str
    client_id: str
    scope: Optional[List[str]] = None

class MFARequest(BaseModel):
    challenge_id: str
    response: str

class TokenRequest(BaseModel):
    grant_type: str
    code: Optional[str] = None
    client_id: str
    client_secret: Optional[str] = None
    redirect_uri: Optional[str] = None
    refresh_token: Optional[str] = None

# 路由
@app.post("/login", tags=["认证"])
async def login(request: LoginRequest):
    """用户登录"""
    logger.info(f"Login request from client: {request.client_id} for user: {request.username}")
    
    # 检查登录失败尝试次数
    if not auth_protection.check_failed_attempts(request.username):
        raise HTTPException(status_code=403, detail="账号已被锁定，请稍后再试")
    
    # 验证用户名和密码
    # 这里简化处理，实际应调用用户管理服务验证
    if request.username != "test" or request.password != "Test@1234":
        # 记录登录失败尝试
        auth_protection.record_failed_attempt(request.username, "127.0.0.1")
        raise HTTPException(status_code=401, detail="用户名或密码错误")
    
    # 记录登录成功
    auth_protection.record_successful_login(request.username, "127.0.0.1")
    
    # 生成多因素认证挑战
    challenge = mfa_service.generate_challenge(
        user_id=request.username,
        mfa_type=MFAType.SMS
    )
    
    # 发送多因素认证挑战
    mfa_service.send_challenge(challenge)
    
    return {
        "challenge_id": challenge.challenge_id,
        "mfa_type": challenge.mfa_type.value,
        "message": "请输入短信验证码"
    }

@app.post("/mfa/verify", tags=["多因素认证"])
async def verify_mfa(request: MFARequest):
    """验证多因素认证"""
    logger.info(f"MFA verification request: {request.challenge_id}")
    
    # 验证多因素认证
    if not mfa_service.verify_challenge(request.challenge_id, request.response):
        raise HTTPException(status_code=401, detail="多因素认证失败")
    
    # 生成 ID Token 和 Access Token
    tokens = oidc_service.generate_tokens(
        user_id="test",
        client_id="mcp-client-001",
        scopes=["openid", "profile", "email"]
    )
    
    return tokens

@app.post("/token", tags=["OIDC"])
async def token(request: TokenRequest):
    """OIDC Token 端点"""
    logger.info(f"Token request: {request.grant_type} from client: {request.client_id}")
    
    # 实现不同 grant_type 的处理逻辑
    if request.grant_type == "authorization_code":
        # 实现授权码模式处理逻辑
        pass
    elif request.grant_type == "refresh_token":
        # 实现刷新令牌模式处理逻辑
        pass
    elif request.grant_type == "client_credentials":
        # 实现客户端凭证模式处理逻辑
        pass
    else:
        raise HTTPException(status_code=400, detail="不支持的 grant_type")
    
    return {}

@app.get("/.well-known/openid-configuration", tags=["OIDC"])
async def discovery_config():
    """OIDC 发现配置"""
    return oidc_service.get_discovery_config()

@app.get("/.well-known/jwks.json", tags=["OIDC"])
async def jwks():
    """JWKS 端点"""
    return oidc_service.get_jwks()

@app.get("/userinfo", tags=["OIDC"])
async def userinfo(access_token: str):
    """UserInfo 端点"""
    return oidc_service.get_user_info(access_token)

@app.post("/password/validate", tags=["密码管理"])
async def validate_password(password: str):
    """验证密码强度"""
    return auth_protection.validate_password_strength(password)

# 运行应用
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)

# 安装依赖
pip install fastapi uvicorn pyjwt pyotp bcrypt

# 配置认证服务
cp config.example.yaml config.yaml
# 编辑配置文件
vim config.yaml

# 启动认证服务
uvicorn mcp_auth_api:app --host 0.0.0.0 --port 8000 --workers 4

# 启动多因素认证服务
uvicorn mcp_mfa_service:app --host 0.0.0.0 --port 8001 --workers 2

# 启动 OIDC 服务
uvicorn mcp_oidc_service:app --host 0.0.0.0 --port 8002 --workers 2