xingrin/backend/apps/scan/notifications/services.py

"""通知服务 - 支持数据库存储和 WebSocket 实时推送"""

import logging
import time
import requests
import urllib3
from .models import Notification, NotificationSettings
from .types import NotificationLevel, NotificationCategory
from .repositories import DjangoNotificationRepository, NotificationSettingsRepository

# 禁用自签名证书的 SSL 警告（远程 Worker 回调场景）
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

logger = logging.getLogger(__name__)


# ============================================================
# 外部推送渠道抽象
# ============================================================

# Discord Embed 颜色映射（使用字符串 key，因为 model 字段存储的是字符串）
DISCORD_COLORS = {
    'low': 0x3498db,       # 蓝色
    'medium': 0xf39c12,    # 橙色
    'high': 0xe74c3c,      # 红色
    'critical': 0x9b59b6,  # 紫色
}

# 分类 emoji（使用字符串 key）
CATEGORY_EMOJI = {
    'scan': '🔍',
    'vulnerability': '⚠️',
    'asset': '🌐',
    'system': '⚙️',
}


def push_to_external_channels(notification: Notification) -> None:
    """
    推送通知到外部渠道（Discord、Slack 等）
    
    根据用户设置决定推送到哪些渠道。
    目前支持：Discord
    未来可扩展：Slack、Telegram、Email 等
    
    Args:
        notification: 通知对象
    """
    settings = NotificationSettings.get_instance()
    
    # 检查分类是否启用
    if not settings.is_category_enabled(notification.category):
        logger.debug(f"分类 {notification.category} 未启用外部推送")
        return
    
    # Discord 渠道
    if settings.discord_enabled and settings.discord_webhook_url:
        try:
            _send_discord(notification, settings.discord_webhook_url)
        except Exception as e:
            logger.warning(f"Discord 推送失败: {e}")
    
    # 未来扩展：Slack
    # if settings.slack_enabled and settings.slack_webhook_url:
    #     _send_slack(notification, settings.slack_webhook_url)
    
    # 未来扩展：Telegram
    # if settings.telegram_enabled and settings.telegram_bot_token:
    #     _send_telegram(notification, settings.telegram_chat_id)


def _send_discord(notification: Notification, webhook_url: str) -> bool:
    """发送到 Discord Webhook"""
    try:
        color = DISCORD_COLORS.get(notification.level, 0x95a5a6)
        emoji = CATEGORY_EMOJI.get(notification.category, '📢')
        
        embed = {
            'title': f"{emoji} {notification.title}",
            'description': notification.message,
            'color': color,
            'footer': {
                'text': f"级别: {notification.get_level_display()} | 分类: {notification.get_category_display()}"
            },
            'timestamp': notification.created_at.isoformat(),
        }
        
        response = requests.post(
            webhook_url,
            json={'embeds': [embed]},
            timeout=10
        )
        
        if response.status_code in (200, 204):
            logger.info(f"Discord 通知发送成功 - {notification.title}")
            return True
        else:
            logger.warning(f"Discord 发送失败 - 状态码: {response.status_code}")
            return False
            
    except requests.RequestException as e:
        logger.error(f"Discord 网络错误: {e}")
        return False


# ============================================================
# 设置服务
# ============================================================

class NotificationSettingsService:
    """通知设置服务"""
    
    def __init__(self, repository: NotificationSettingsRepository | None = None):
        self.repo = repository or NotificationSettingsRepository()
    
    def get_settings(self) -> dict:
        """获取通知设置（前端格式）"""
        settings = self.repo.get_settings()
        return {
            'discord': {
                'enabled': settings.discord_enabled,
                'webhookUrl': settings.discord_webhook_url,
            },
            'categories': settings.categories,
        }
    
    def update_settings(self, data: dict) -> dict:
        """更新通知设置
        
        注意：DRF CamelCaseJSONParser 会将前端的 webhookUrl 转换为 webhook_url
        """
        discord_data = data.get('discord', {})
        categories = data.get('categories', {})
        
        # CamelCaseJSONParser 转换后的字段名是 webhook_url
        webhook_url = discord_data.get('webhook_url', '')
        
        settings = self.repo.update_settings(
            discord_enabled=discord_data.get('enabled', False),
            discord_webhook_url=webhook_url,
            categories=categories,
        )
        
        return {
            'discord': {
                'enabled': settings.discord_enabled,
                'webhookUrl': settings.discord_webhook_url,
            },
            'categories': settings.categories,
        }


class NotificationService:
    """通知业务服务，封装常用查询与更新操作"""

    def __init__(self, repository: DjangoNotificationRepository | None = None):
        self.repo = repository or DjangoNotificationRepository()

    def get_notifications(self, level: str | None = None, unread: bool | None = None):
        return self.repo.get_filtered(level=level, unread=unread)

    def get_unread_count(self) -> int:
        return self.repo.get_unread_count()

    def mark_all_as_read(self) -> int:
        return self.repo.mark_all_as_read()


def create_notification(
    title: str,
    message: str,
    level: NotificationLevel = NotificationLevel.LOW,
    category: NotificationCategory = NotificationCategory.SYSTEM
) -> Notification:
    """
    创建通知记录并实时推送
    
    增强的重试机制：
    - 最多重试 3 次
    - 每次重试前强制关闭并重建数据库连接
    - 重试间隔：1秒 → 2秒 → 3秒
    - 针对连接错误进行特殊处理
    
    Args:
        title: 通知标题
        message: 通知消息
        level: 通知级别
        category: 通知分类
        
    Returns:
        Notification: 创建的通知对象
        
    Raises:
        Exception: 重试3次后仍然失败
    """
    from django.db import connection
    from psycopg2 import OperationalError, InterfaceError

    repo = DjangoNotificationRepository()

    max_retries = 3
    last_exception = None
    
    for attempt in range(1, max_retries + 1):
        try:
            # 强制关闭旧连接并重建（每次尝试都重建）
            if attempt > 1:
                logger.debug(f"重试创建通知 ({attempt}/{max_retries}) - {title}")
            
            connection.close()
            connection.ensure_connection()
            
            # 测试连接是否真的可用
            connection.cursor().execute("SELECT 1")
            
            # 1. 写入数据库（通过仓储层统一访问 ORM）
            notification = repo.create(
                title=title,
                message=message,
                level=level,
                category=category,
            )
            
            # 2. WebSocket 实时推送（推送失败不影响通知创建）
            try:
                _push_to_websocket(notification)
            except Exception as push_error:
                logger.warning(f"WebSocket 推送失败，但通知已创建 - {title}: {push_error}")
            
            # 3. 外部渠道推送（Discord/Slack 等，推送失败不影响通知创建）
            try:
                push_to_external_channels(notification)
            except Exception as external_error:
                logger.warning(f"外部渠道推送失败，但通知已创建 - {title}: {external_error}")
            
            if attempt > 1:
                logger.info(f"✓ 通知创建成功（重试 {attempt-1} 次后） - {title}")
            else:
                logger.debug(f"通知已创建并推送 - {title}")
            
            return notification
            
        except (OperationalError, InterfaceError) as e:
            # 数据库连接错误，需要重试
            last_exception = e
            error_msg = str(e)
            logger.warning(
                f"数据库连接错误 ({attempt}/{max_retries}) - {title}: {error_msg[:100]}"
            )
            
            if attempt < max_retries:
                # 指数退避：1秒、2秒、3秒
                sleep_time = attempt
                logger.debug(f"等待 {sleep_time} 秒后重试...")
                time.sleep(sleep_time)
            else:
                logger.error(
                    f"创建通知失败 - 数据库连接问题（已重试 {max_retries} 次） - {title}: {error_msg}"
                )
                
        except Exception as e:
            # 其他错误，不重试直接抛出
            last_exception = e
            error_str = str(e).lower()
            
            if 'connection' in error_str or 'closed' in error_str:
                logger.error(f"创建通知失败 - 连接相关错误 - {title}: {e}")
            else:
                logger.error(f"创建通知失败 - {title}: {e}")
            
            # 非连接错误，直接抛出不重试
            raise
    
    # 所有重试都失败了
    error_msg = f"创建通知失败 - 已重试 {max_retries} 次仍然失败 - {title}"
    logger.error(error_msg)
    raise RuntimeError(error_msg) from last_exception


def _push_to_websocket(notification: Notification) -> None:
    """
    推送通知到 WebSocket 客户端
    
    - 在 Server 容器中：直接通过 Channel Layer 推送
    - 在 Worker 容器中：通过 API 回调让 Server 推送（因为 Worker 无法访问 Redis）
    """
    import os
    
    # 检测是否在 Worker 容器中（有 SERVER_URL 环境变量）
    server_url = os.environ.get("SERVER_URL")
    
    if server_url:
        # Worker 容器：通过 API 回调
        _push_via_api_callback(notification, server_url)
    else:
        # Server 容器：直接推送
        _push_via_channel_layer(notification)


def _push_via_api_callback(notification: Notification, server_url: str) -> None:
    """
    通过 HTTP 回调推送通知（Worker → Server 跨容器通信）
    
    注意：这不是同进程内的 service 调用 view，而是 Worker 容器
    通过 HTTP 请求 Server 容器的 /api/callbacks/notification/ 接口。
    Worker 无法直接访问 Redis，需要由 Server 代为推送 WebSocket。
    """
    import requests
    
    try:
        callback_url = f"{server_url}/api/callbacks/notification/"
        data = {
            'id': notification.id,
            'category': notification.category,
            'title': notification.title,
            'message': notification.message,
            'level': notification.level,
            'created_at': notification.created_at.isoformat()
        }
        
        # verify=False: 远程 Worker 回调 Server 时可能使用自签名证书
        resp = requests.post(callback_url, json=data, timeout=5, verify=False)
        resp.raise_for_status()
        
        logger.debug(f"通知回调推送成功 - ID: {notification.id}")
        
    except Exception as e:
        logger.warning(f"通知回调推送失败 - ID: {notification.id}: {e}")


def _push_via_channel_layer(notification: Notification) -> None:
    """通过 Channel Layer 直接推送通知（Server 容器使用）"""
    try:
        logger.debug(f"开始推送通知到 WebSocket - ID: {notification.id}")
        
        from asgiref.sync import async_to_sync
        from channels.layers import get_channel_layer
        
        # 获取 Channel Layer
        channel_layer = get_channel_layer()
        
        if channel_layer is None:
            logger.warning("Channel Layer 未配置，跳过 WebSocket 推送")
            return
        
        # 构造通知数据
        data = {
            'type': 'notification.message',  # 对应 Consumer 的 notification_message 方法
            'id': notification.id,
            'category': notification.category,
            'title': notification.title,
            'message': notification.message,
            'level': notification.level,
            'created_at': notification.created_at.isoformat()
        }
        
        # 发送到通知组（所有连接的客户端）
        async_to_sync(channel_layer.group_send)(
            'notifications',  # 组名
            data
        )
        
        logger.debug(f"通知推送成功 - ID: {notification.id}")
        
    except ImportError as e:
        logger.warning(f"Channels 模块未安装，跳过 WebSocket 推送: {e}")
    except Exception as e:
        logger.warning(f"WebSocket 推送失败 - ID: {notification.id}: {e}", exc_info=True)