ipam/backend/app/services/user_service.py

from typing import Optional, List
from sqlalchemy.orm import Session
from datetime import datetime, timedelta
from fastapi import HTTPException, status

from app.models.auth import User, UserRole, UserStatus, Permission, RolePermission, RefreshToken
from app.core.security import SecurityUtils


# 默认权限项 + 角色权限映射
DEFAULT_PERMISSIONS = [
    # 网段
    ("network:view",    "查看网段", "network"),
    ("network:create",  "创建网段", "network"),
    ("network:edit",    "编辑网段", "network"),
    ("network:delete",  "删除网段", "network"),
    # IP
    ("ip:view",         "查看 IP", "ip"),
    ("ip:edit",         "编辑 IP", "ip"),
    # 扫描
    ("scan:run",        "执行扫描", "scan"),
    ("scan:view",       "查看扫描结果", "scan"),
    # SNMP
    ("snmp:view",       "查看 SNMP", "snmp"),
    ("snmp:edit",       "配置 SNMP", "snmp"),
    # 告警
    ("alert:view",      "查看告警", "alert"),
    ("alert:ack",       "确认告警", "alert"),
    ("alert:resolve",   "解决告警", "alert"),
    # 报告
    ("report:export",   "导出报告", "report"),
    # 用户管理
    ("user:view",       "查看用户", "system"),
    ("user:create",     "创建用户", "system"),
    ("user:edit",       "编辑用户", "system"),
    ("user:delete",     "删除用户", "system"),
    # 审计
    ("audit:view",      "查看审计日志", "system"),
]


# 角色 -> 权限代码集合
DEFAULT_ROLE_PERMISSIONS = {
    UserRole.SUPER_ADMIN: {p[0] for p in DEFAULT_PERMISSIONS},  # 全部
    UserRole.ADMIN: {
        "network:view", "network:create", "network:edit", "network:delete",
        "ip:view", "ip:edit",
        "scan:run", "scan:view",
        "snmp:view", "snmp:edit",
        "alert:view", "alert:ack", "alert:resolve",
        "report:export",
        "user:view", "audit:view",
    },
    UserRole.OPERATOR: {
        "network:view",
        "ip:view", "ip:edit",
        "scan:run", "scan:view",
        "snmp:view",
        "alert:view", "alert:ack",
        "report:export",
    },
    UserRole.VIEWER: {
        "network:view",
        "ip:view",
        "scan:view",
        "snmp:view",
        "alert:view",
    },
}


class UserService:
    """用户服务"""

    @staticmethod
    def init_default_permissions(db: Session) -> int:
        """初始化默认权限项和角色权限映射（已存在则跳过）"""
        created = 0
        perm_map = {}
        for code, name, category in DEFAULT_PERMISSIONS:
            existing = db.query(Permission).filter(Permission.code == code).first()
            if existing:
                perm_map[code] = existing
                continue
            p = Permission(code=code, name=name, category=category)
            db.add(p)
            db.flush()
            perm_map[code] = p
            created += 1

        # 角色 -> 权限关联
        for role, codes in DEFAULT_ROLE_PERMISSIONS.items():
            for code in codes:
                if code not in perm_map:
                    continue
                pid = perm_map[code].id
                link_existing = db.query(RolePermission).filter(
                    RolePermission.role == role,
                    RolePermission.permission_id == pid
                ).first()
                if link_existing:
                    continue
                db.add(RolePermission(role=role, permission_id=pid))

        db.commit()
        return created

    @staticmethod
    def get_user_permissions(db: Session, user: User) -> set:
        """获取用户拥有的全部权限代码集合"""
        if user.role == UserRole.SUPER_ADMIN:
            return {p.code for p in db.query(Permission).all()}
        # 其他角色查 role_permissions 关联表
        rows = db.query(Permission.code).join(
            RolePermission, RolePermission.permission_id == Permission.id
        ).filter(RolePermission.role == user.role).all()
        return {r[0] for r in rows}
    
    @staticmethod
    def init_default_admin(db: Session):
        """初始化默认管理员账户"""
        admin_exists = db.query(User).filter(User.username == "admin").first()
        if not admin_exists:
            admin = User(
                username="admin",
                email="admin@ipam.local",
                real_name="系统管理员",
                password_hash=SecurityUtils.hash_password("admin123"),
                role=UserRole.SUPER_ADMIN,
                status=UserStatus.ACTIVE
            )
            db.add(admin)
            db.commit()
            db.refresh(admin)
            return admin
        return None
    
    @staticmethod
    def authenticate(db: Session, username: str, password: str) -> Optional[User]:
        """用户认证"""
        user = db.query(User).filter(User.username == username).first()
        
        if not user:
            return None
        
        if not SecurityUtils.verify_password(password, user.password_hash):
            # 记录登录失败
            user.login_failed_count += 1
            # 登录失败5次锁定账号
            if user.login_failed_count >= 5:
                user.status = UserStatus.LOCKED
            db.commit()
            return None
        
        # 登录成功，重置失败计数
        if user.login_failed_count > 0:
            user.login_failed_count = 0
            db.commit()
        
        # 检查用户状态
        if user.status != UserStatus.ACTIVE:
            return None
        
        return user
    
    @staticmethod
    def create_user(db: Session,
                    username: str,
                    password: str,
                    email: Optional[str] = None,
                    phone: Optional[str] = None,
                    real_name: Optional[str] = None,
                    role: UserRole = UserRole.VIEWER,
                    created_by: Optional[int] = None) -> User:
        """创建用户"""
        # 检查用户名是否存在
        existing = db.query(User).filter(User.username == username).first()
        if existing:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail="用户名已存在"
            )
        
        # 检查邮箱是否存在
        if email:
            existing_email = db.query(User).filter(User.email == email).first()
            if existing_email:
                raise HTTPException(
                    status_code=status.HTTP_400_BAD_REQUEST,
                    detail="邮箱已被使用"
                )
        
        user = User(
            username=username,
            email=email,
            phone=phone,
            real_name=real_name,
            password_hash=SecurityUtils.hash_password(password),
            role=role,
            status=UserStatus.ACTIVE,
            created_by=created_by
        )
        
        db.add(user)
        db.commit()
        db.refresh(user)
        return user
    
    @staticmethod
    def update_login_info(db: Session, user_id: int, ip_address: str = None):
        """更新登录信息"""
        user = db.query(User).filter(User.id == user_id).first()
        if user:
            user.last_login_at = datetime.utcnow()
            user.last_login_ip = ip_address
            db.commit()
    
    @staticmethod
    def get_user_by_id(db: Session, user_id: int) -> Optional[User]:
        """根据ID获取用户"""
        return db.query(User).filter(User.id == user_id).first()
    
    @staticmethod
    def get_user_by_username(db: Session, username: str) -> Optional[User]:
        """根据用户名获取用户"""
        return db.query(User).filter(User.username == username).first()
    
    @staticmethod
    def get_users(db: Session, skip: int = 0, limit: int = 50, status: UserStatus = None, role: UserRole = None) -> tuple[int, List[User]]:
        """获取用户列表"""
        query = db.query(User)
        
        if status:
            query = query.filter(User.status == status)
        if role:
            query = query.filter(User.role == role)
        
        total = query.count()
        items = query.order_by(User.id.desc()).offset(skip).limit(limit).all()
        
        return total, items
    
    @staticmethod
    def change_password(db: Session, user_id: int, old_password: str, new_password: str) -> bool:
        """修改密码"""
        user = db.query(User).filter(User.id == user_id).first()
        if not user:
            return False
        
        if not SecurityUtils.verify_password(old_password, user.password_hash):
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail="原密码不正确"
            )
        
        user.password_hash = SecurityUtils.hash_password(new_password)
        db.commit()
        return True
    
    @staticmethod
    def reset_password(db: Session, user_id: int, new_password: str) -> bool:
        """管理员重置密码"""
        user = db.query(User).filter(User.id == user_id).first()
        if not user:
            return False
        
        user.password_hash = SecurityUtils.hash_password(new_password)
        db.commit()
        return True
    
    @staticmethod
    def update_user_status(db: Session, user_id: int, status: UserStatus) -> Optional[User]:
        """更新用户状态"""
        user = db.query(User).filter(User.id == user_id).first()
        if not user:
            return None
        
        user.status = status
        db.commit()
        db.refresh(user)
        return user
    
    @staticmethod
    def delete_user(db: Session, user_id: int) -> bool:
        """删除用户（软删除，实际标记为禁用）"""
        user = db.query(User).filter(User.id == user_id).first()
        if not user:
            return False
        
        user.status = UserStatus.INACTIVE
        db.commit()
        return True
    
    @staticmethod
    def unlock_user(db: Session, user_id: int) -> Optional[User]:
        """解锁被锁定的用户"""
        user = db.query(User).filter(User.id == user_id).first()
        if not user:
            return None
        
        if user.status == UserStatus.LOCKED:
            user.status = UserStatus.ACTIVE
            user.login_failed_count = 0
            db.commit()
            db.refresh(user)
        
        return user


class TokenService:
    """令牌服务"""
    
    @staticmethod
    def create_tokens(db: Session, user: User, ip_address: str = None, user_agent: str = None) -> dict:
        """创建访问令牌和刷新令牌"""
        access_token_expires = timedelta(minutes=120)  # 访问令牌2小时
        access_token = SecurityUtils.create_access_token(
            data={
                "sub": user.username,
                "user_id": user.id,
                "role": user.role.value
            },
            expires_delta=access_token_expires
        )
        
        refresh_token_expires = timedelta(days=7)  # 刷新令牌7天
        refresh_token_value = SecurityUtils.create_refresh_token(
            data={
                "sub": user.username,
                "user_id": user.id
            },
            expires_delta=refresh_token_expires
        )
        
        # 保存刷新令牌
        refresh_token = RefreshToken(
            user_id=user.id,
            token=refresh_token_value,
            expires_at=datetime.utcnow() + refresh_token_expires,
            ip_address=ip_address,
            user_agent=user_agent
        )
        db.add(refresh_token)
        db.commit()
        
        return {
            "access_token": access_token,
            "refresh_token": refresh_token_value,
            "token_type": "bearer",
            "expires_in": int(access_token_expires.total_seconds())
        }
    
    @staticmethod
    def refresh_access_token(db: Session, refresh_token: str) -> Optional[dict]:
        """使用刷新令牌获取新的访问令牌"""
        try:
            payload = SecurityUtils.verify_token(refresh_token)
            username = payload.get("sub")
            user_id = payload.get("user_id")
            
            # 检查刷新令牌是否有效且未被撤销
            token_record = db.query(RefreshToken).filter(
                RefreshToken.token == refresh_token,
                RefreshToken.user_id == user_id,
                RefreshToken.revoked == False,
                RefreshToken.expires_at > datetime.utcnow()
            ).first()
            
            if not token_record:
                return None
            
            user = db.query(User).filter(User.id == user_id).first()
            if not user or user.status != UserStatus.ACTIVE:
                return None
            
            # 创建新的访问令牌
            access_token_expires = timedelta(minutes=120)
            access_token = SecurityUtils.create_access_token(
                data={
                    "sub": user.username,
                    "user_id": user.id,
                    "role": user.role.value
                },
                expires_delta=access_token_expires
            )
            
            return {
                "access_token": access_token,
                "token_type": "bearer",
                "expires_in": int(access_token_expires.total_seconds())
            }
            
        except Exception:
            return None
    
    @staticmethod
    def revoke_token(db: Session, refresh_token: str) -> bool:
        """撤销刷新令牌"""
        token_record = db.query(RefreshToken).filter(RefreshToken.token == refresh_token).first()
        if token_record:
            token_record.revoked = True
            token_record.revoked_at = datetime.utcnow()
            db.commit()
            return True
        return False
    
    @staticmethod
    def revoke_all_user_tokens(db: Session, user_id: int) -> int:
        """撤销用户所有刷新令牌"""
        tokens = db.query(RefreshToken).filter(
            RefreshToken.user_id == user_id,
            RefreshToken.revoked == False
        ).all()
        
        for token in tokens:
            token.revoked = True
            token.revoked_at = datetime.utcnow()
        
        db.commit()
        return len(tokens)