增加企业微信

Author: yyhuni

backend/apps/scan/migrations/0003_add_wecom_fields.py

@@ -0,0 +1,23 @@
+# Generated manually for WeCom notification support
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('scan', '0002_add_cached_screenshots_count'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='notificationsettings',
+            name='wecom_enabled',
+            field=models.BooleanField(default=False, help_text='是否启用企业微信通知'),
+        ),
+        migrations.AddField(
+            model_name='notificationsettings',
+            name='wecom_webhook_url',
+            field=models.URLField(blank=True, default='', help_text='企业微信机器人 Webhook URL'),
+        ),
+    ]

backend/apps/scan/notifications/models.py

@@ -1,40 +1,49 @@
 """通知系统数据模型"""
 
+import logging
+from datetime import timedelta
+
 from django.db import models
+from django.utils import timezone
+
+from .types import NotificationCategory, NotificationLevel
 
-from .types import NotificationLevel, NotificationCategory
+logger = logging.getLogger(__name__)
 
 
 class NotificationSettings(models.Model):
     """
     通知设置（单例模型）
     存储 Discord webhook 配置和各分类的通知开关
     """
-    
+
     # Discord 配置
     discord_enabled = models.BooleanField(default=False, help_text='是否启用 Discord 通知')
     discord_webhook_url = models.URLField(blank=True, default='', help_text='Discord Webhook URL')
-    
+
+    # 企业微信配置
+    wecom_enabled = models.BooleanField(default=False, help_text='是否启用企业微信通知')
+    wecom_webhook_url = models.URLField(blank=True, default='', help_text='企业微信机器人 Webhook URL')
+
     # 分类开关（使用 JSONField 存储）
     categories = models.JSONField(
         default=dict,
         help_text='各分类通知开关，如 {"scan": true, "vulnerability": true, "asset": true, "system": false}'
     )
-    
+
     # 时间信息
     created_at = models.DateTimeField(auto_now_add=True)
     updated_at = models.DateTimeField(auto_now=True)
-    
+
     class Meta:
         db_table = 'notification_settings'
         verbose_name = '通知设置'
         verbose_name_plural = '通知设置'
-    
+
     def save(self, *args, **kwargs):
-        # 单例模式：强制只有一条记录
-        self.pk = 1
+        self.pk = 1  # 单例模式
         super().save(*args, **kwargs)
-    
+
     @classmethod
     def get_instance(cls) -> 'NotificationSettings':
         """获取或创建单例实例"""
@@ -52,44 +61,41 @@ def get_instance(cls) -> 'NotificationSettings':
             }
         )
         return obj
-    
+
     def is_category_enabled(self, category: str) -> bool:
         """检查指定分类是否启用通知"""
         return self.categories.get(category, False)
 
 
 class Notification(models.Model):
     """通知模型"""
-    
+
     id = models.AutoField(primary_key=True)
-    
-    # 通知分类
+
     category = models.CharField(
         max_length=20,
         choices=NotificationCategory.choices,
         default=NotificationCategory.SYSTEM,
         db_index=True,
         help_text='通知分类'
     )
-    
-    # 通知级别
+
     level = models.CharField(
         max_length=20,
         choices=NotificationLevel.choices,
         default=NotificationLevel.LOW,
         db_index=True,
         help_text='通知级别'
     )
-    
+
     title = models.CharField(max_length=200, help_text='通知标题')
     message = models.CharField(max_length=2000, help_text='通知内容')
-    
-    # 时间信息
+
     created_at = models.DateTimeField(auto_now_add=True, help_text='创建时间')
-    
+
     is_read = models.BooleanField(default=False, help_text='是否已读')
     read_at = models.DateTimeField(null=True, blank=True, help_text='阅读时间')
-    
+
     class Meta:
         db_table = 'notification'
         verbose_name = '通知'
@@ -101,44 +107,26 @@ class Meta:
             models.Index(fields=['level', '-created_at']),
             models.Index(fields=['is_read', '-created_at']),
         ]
-    
+
     def __str__(self):
         return f"{self.get_level_display()} - {self.title}"
-    
+
     @classmethod
-    def cleanup_old_notifications(cls):
-        """
-        清理超过15天的旧通知（硬编码）
-        
-        Returns:
-            int: 删除的通知数量
-        """
-        from datetime import timedelta
-        from django.utils import timezone
-        
-        # 硬编码：只保留最近15天的通知
+    def cleanup_old_notifications(cls) -> int:
+        """清理超过15天的旧通知"""
         cutoff_date = timezone.now() - timedelta(days=15)
-        delete_result = cls.objects.filter(created_at__lt=cutoff_date).delete()
-        
-        return delete_result[0] if delete_result[0] else 0
-    
+        deleted_count, _ = cls.objects.filter(created_at__lt=cutoff_date).delete()
+        return deleted_count or 0
+
     def save(self, *args, **kwargs):
-        """
-        重写save方法，在创建新通知时自动清理旧通知
-        """
+        """重写save方法，在创建新通知时自动清理旧通知"""
         is_new = self.pk is None
         super().save(*args, **kwargs)
-        
-        # 只在创建新通知时执行清理（自动清理超过15天的通知）
+
         if is_new:
             try:
                 deleted_count = self.__class__.cleanup_old_notifications()
                 if deleted_count > 0:
-                    import logging
-                    logger = logging.getLogger(__name__)
-                    logger.info(f"自动清理了 {deleted_count} 条超过15天的旧通知")
-            except Exception as e:
-                # 清理失败不应影响通知创建
-                import logging
-                logger = logging.getLogger(__name__)
-                logger.warning(f"通知自动清理失败: {e}")
+                    logger.info("自动清理了 %d 条超过15天的旧通知", deleted_count)
+            except Exception:
+                logger.warning("通知自动清理失败", exc_info=True)

backend/apps/scan/notifications/repositories.py

@@ -1,52 +1,70 @@
+"""通知系统仓储层模块"""
+
 import logging
-from typing import TypedDict
+from dataclasses import dataclass
+from typing import Optional
+
+from django.db.models import QuerySet
 from django.utils import timezone
 
 from apps.common.decorators import auto_ensure_db_connection
-from .models import Notification, NotificationSettings
 
+from .models import Notification, NotificationSettings
 
 logger = logging.getLogger(__name__)
 
 
-class NotificationSettingsData(TypedDict):
-    """通知设置数据结构"""
+@dataclass
+class NotificationSettingsData:
+    """通知设置更新数据"""
+
     discord_enabled: bool
     discord_webhook_url: str
     categories: dict[str, bool]
+    wecom_enabled: bool = False
+    wecom_webhook_url: str = ''
 
 
 @auto_ensure_db_connection
 class NotificationSettingsRepository:
     """通知设置仓储层"""
-    
+
     def get_settings(self) -> NotificationSettings:
         """获取通知设置单例"""
         return NotificationSettings.get_instance()
-    
-    def update_settings(
-        self,
-        discord_enabled: bool,
-        discord_webhook_url: str,
-        categories: dict[str, bool]
-    ) -> NotificationSettings:
+
+    def update_settings(self, data: NotificationSettingsData) -> NotificationSettings:
         """更新通知设置"""
         settings = NotificationSettings.get_instance()
-        settings.discord_enabled = discord_enabled
-        settings.discord_webhook_url = discord_webhook_url
-        settings.categories = categories
+        settings.discord_enabled = data.discord_enabled
+        settings.discord_webhook_url = data.discord_webhook_url
+        settings.wecom_enabled = data.wecom_enabled
+        settings.wecom_webhook_url = data.wecom_webhook_url
+        settings.categories = data.categories
         settings.save()
         return settings
-    
+
     def is_category_enabled(self, category: str) -> bool:
         """检查指定分类是否启用"""
-        settings = self.get_settings()
-        return settings.is_category_enabled(category)
+        return self.get_settings().is_category_enabled(category)
 
 
 @auto_ensure_db_connection
 class DjangoNotificationRepository:
-    def get_filtered(self, level: str | None = None, unread: bool | None = None):
+    """通知数据仓储层"""
+
+    def get_filtered(
+        self,
+        level: Optional[str] = None,
+        unread: Optional[bool] = None
+    ) -> QuerySet[Notification]:
+        """
+        获取过滤后的通知列表
+
+        Args:
+            level: 通知级别过滤
+            unread: 已读状态过滤 (True=未读, False=已读, None=全部)
+        """
         queryset = Notification.objects.all()
 
         if level:
@@ -60,16 +78,24 @@ def get_filtered(self, level: str | None = None, unread: bool | None = None):
         return queryset.order_by("-created_at")
 
     def get_unread_count(self) -> int:
+        """获取未读通知数量"""
         return Notification.objects.filter(is_read=False).count()
 
     def mark_all_as_read(self) -> int:
-        updated = Notification.objects.filter(is_read=False).update(
+        """标记所有通知为已读，返回更新数量"""
+        return Notification.objects.filter(is_read=False).update(
             is_read=True,
             read_at=timezone.now(),
         )
-        return updated
 
-    def create(self, title: str, message: str, level: str, category: str = 'system') -> Notification:
+    def create(
+        self,
+        title: str,
+        message: str,
+        level: str,
+        category: str = 'system'
+    ) -> Notification:
+        """创建新通知"""
         return Notification.objects.create(
             category=category,
             level=level,