SenSu/services/auth_service.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import logging
import os
import hashlib
import secrets
from typing import Dict, List, Optional
from dataclasses import dataclass
import time
import json

logger = logging.getLogger(__name__)

@dataclass
class User:
    """用户数据类"""
    username: str
    password_hash: str
    permissions: List[str]
    is_active: bool = True
    created_at: float = None

@dataclass
class Token:
    """令牌数据类"""
    token: str
    username: str
    permissions: List[str]
    created_at: float
    expires_at: float
    is_valid: bool = True

class AuthService:
    """认证服务 - 处理用户认证和权限验证"""
    
    def __init__(self, config: Dict):
        self.config = config
        self.users: Dict[str, User] = {}
        self.tokens: Dict[str, Token] = {}
        self.token_expiry_hours = 24
        self.secret_key = secrets.token_hex(32)
        logger.debug("AuthService初始化开始")
        
        # 初始化默认用户
        self._init_default_users()
    
    def _init_default_users(self):
        """初始化默认用户"""
        try:
            # 创建默认管理员用户
            admin_password_hash = self._hash_password(os.environ.get("SENSU_ADMIN_PASSWORD","admin123"))
            admin_user = User(
                username="admin",
                password_hash=admin_password_hash,
                permissions=["admin"],
                created_at=time.time()
            )
            self.users["admin"] = admin_user
            
            # 创建默认API用户
            api_password_hash = self._hash_password(os.environ.get("SENSU_API_PASSWORD","api123"))
            api_user = User(
                username="api",
                password_hash=api_password_hash,
                permissions=["framework.status.read", "plugin.info.read"],
                created_at=time.time()
            )
            self.users["api"] = api_user
            
            logger.debug("默认用户初始化完成")
            
        except Exception as e:
            logger.error(f"初始化默认用户时出错: {str(e)}", exc_info=True)
            raise
    
    def _hash_password(self, password: str) -> str:
        """哈希密码"""
        try:
            salt = "catframework_salt"  # 实际应该使用随机盐
            return hashlib.sha256((password + salt).encode()).hexdigest()
        except Exception as e:
            logger.error(f"哈希密码时出错: {str(e)}", exc_info=True)
            raise
    
    def authenticate_user(self, username: str, password: str) -> Optional[Token]:
        """用户认证"""
        try:
            logger.debug(f"用户认证尝试: {username}")
            
            if username not in self.users:
                logger.warning(f"用户不存在: {username}")
                return None
            
            user = self.users[username]
            
            if not user.is_active:
                logger.warning(f"用户已被禁用: {username}")
                return None
            
            password_hash = self._hash_password(password)
            if user.password_hash != password_hash:
                logger.warning(f"密码错误: {username}")
                return None
            
            # 创建令牌
            token = self._create_token(user)
            logger.debug(f"用户认证成功: {username}")
            return token
            
        except Exception as e:
            logger.error(f"用户认证时出错: {str(e)}", exc_info=True)
            return None
    
    def _create_token(self, user: User) -> Token:
        """创建令牌"""
        try:
            token_str = secrets.token_hex(32)
            created_at = time.time()
            expires_at = created_at + (self.token_expiry_hours * 3600)
            
            token = Token(
                token=token_str,
                username=user.username,
                permissions=user.permissions,
                created_at=created_at,
                expires_at=expires_at
            )
            
            self.tokens[token_str] = token
            logger.debug(f"创建令牌: {user.username}, 有效期: {self.token_expiry_hours}小时")
            return token
            
        except Exception as e:
            logger.error(f"创建令牌时出错: {str(e)}", exc_info=True)
            raise
    
    def validate_token(self, token_str: str) -> Optional[Token]:
        """验证令牌"""
        try:
            if token_str not in self.tokens:
                logger.debug("令牌不存在")
                return None
            
            token = self.tokens[token_str]
            
            # 检查令牌是否有效
            if not token.is_valid:
                logger.debug("令牌已失效")
                return None
            
            # 检查令牌是否过期
            if time.time() > token.expires_at:
                logger.debug("令牌已过期")
                token.is_valid = False
                return None
            
            logger.debug(f"令牌验证成功: {token.username}")
            return token
            
        except Exception as e:
            logger.error(f"验证令牌时出错: {str(e)}", exc_info=True)
            return None
    
    def revoke_token(self, token_str: str) -> bool:
        """撤销令牌"""
        try:
            if token_str in self.tokens:
                self.tokens[token_str].is_valid = False
                logger.debug(f"令牌已撤销: {token_str}")
                return True
            else:
                logger.warning(f"要撤销的令牌不存在: {token_str}")
                return False
        except Exception as e:
            logger.error(f"撤销令牌时出错: {str(e)}", exc_info=True)
            return False
    
    def check_permission(self, token_str: str, permission: str) -> bool:
        """检查权限"""
        try:
            token = self.validate_token(token_str)
            if not token:
                return False
            
            # 检查admin权限
            if "admin" in token.permissions:
                return True
            
            # 检查具体权限
            has_permission = permission in token.permissions
            
            logger.debug(f"权限检查: {token.username} -> {permission} = {has_permission}")
            return has_permission
            
        except Exception as e:
            logger.error(f"检查权限时出错: {str(e)}", exc_info=True)
            return False
    
    def create_user(self, username: str, password: str, permissions: List[str]) -> bool:
        """创建用户"""
        try:
            if username in self.users:
                logger.warning(f"用户已存在: {username}")
                return False
            
            password_hash = self._hash_password(password)
            user = User(
                username=username,
                password_hash=password_hash,
                permissions=permissions,
                created_at=time.time()
            )
            
            self.users[username] = user
            logger.debug(f"用户创建成功: {username}, 权限: {permissions}")
            return True
            
        except Exception as e:
            logger.error(f"创建用户时出错: {str(e)}", exc_info=True)
            return False
    
    def get_user_info(self, username: str) -> Optional[Dict]:
        """获取用户信息"""
        try:
            if username not in self.users:
                return None
            
            user = self.users[username]
            return {
                "username": user.username,
                "permissions": user.permissions,
                "is_active": user.is_active,
                "created_at": user.created_at
            }
        except Exception as e:
            logger.error(f"获取用户信息时出错: {str(e)}", exc_info=True)
            return None
    
    def cleanup_expired_tokens(self):
        """清理过期令牌"""
        try:
            current_time = time.time()
            expired_tokens = []
            
            for token_str, token in self.tokens.items():
                if current_time > token.expires_at:
                    expired_tokens.append(token_str)
            
            for token_str in expired_tokens:
                del self.tokens[token_str]
            
            if expired_tokens:
                logger.debug(f"清理了 {len(expired_tokens)} 个过期令牌")
                
        except Exception as e:
            logger.error(f"清理过期令牌时出错: {str(e)}", exc_info=True)
    
    def shutdown(self):
        """关闭认证服务"""
        try:
            logger.info("关闭认证服务")
            self.cleanup_expired_tokens()
            self.users.clear()
            self.tokens.clear()
            logger.debug("认证服务关闭完成")
        except Exception as e:
            logger.error(f"关闭认证服务时出错: {str(e)}", exc_info=True)