ProxyPool/services/plugin_service.py

"""插件业务服务"""
from datetime import datetime
from typing import List, Optional
from core.db import get_db
from core.plugin_system.registry import registry
from core.plugin_system.base import BaseCrawlerPlugin
from repositories.settings_repo import PluginSettingsRepository
from models.domain import PluginInfo, ProxyRaw
from core.log import logger


class PluginService:
    """插件业务服务：管理插件生命周期、执行爬取、配置管理"""

    def __init__(self):
        self.plugin_settings_repo = PluginSettingsRepository()
        self._stats: dict[str, dict] = {}

    async def list_plugins(self) -> List[PluginInfo]:
        """获取所有插件信息（合并持久化状态和配置）"""
        async with get_db() as db:
            db_states = await self.plugin_settings_repo.list_all(db)

        result = []
        for plugin in registry.list_plugins():
            # 合并持久化状态
            state = db_states.get(plugin.name, {})
            if "enabled" in state:
                plugin.enabled = state["enabled"]
            if "config" in state and isinstance(state["config"], dict):
                plugin.update_config(state["config"])

            stat = self._stats.get(plugin.name, {
                "success_count": 0,
                "failure_count": 0,
                "last_run": None,
            })
            result.append(PluginInfo(
                id=plugin.name,
                name=plugin.name,
                display_name=plugin.display_name or plugin.name,
                description=plugin.description or f"从 {plugin.name} 爬取代理",
                enabled=plugin.enabled,
                last_run=stat.get("last_run"),
                success_count=stat.get("success_count", 0),
                failure_count=stat.get("failure_count", 0),
            ))
        return result

    async def toggle_plugin(self, plugin_id: str, enabled: bool) -> bool:
        plugin = registry.get(plugin_id)
        if not plugin:
            return False
        async with get_db() as db:
            success = await self.plugin_settings_repo.set_enabled(db, plugin_id, enabled)
        if success:
            plugin.enabled = enabled
            logger.info(f"Plugin {plugin_id} toggled to {enabled}")
        return success

    async def get_plugin_config(self, plugin_id: str) -> Optional[dict]:
        """获取插件当前配置（合并默认值和持久化值）"""
        plugin = registry.get(plugin_id)
        if not plugin:
            return None
        async with get_db() as db:
            saved = await self.plugin_settings_repo.get_config(db, plugin_id)
        config = dict(plugin.default_config)
        if saved:
            config.update(saved)
        return config

    async def update_plugin_config(self, plugin_id: str, config: dict) -> bool:
        """更新插件配置（只保存已存在于 default_config 中的键）"""
        plugin = registry.get(plugin_id)
        if not plugin:
            return False
        # 过滤非法键
        safe_config = {k: v for k, v in config.items() if k in plugin.default_config}
        if not safe_config:
            return False
        plugin.update_config(safe_config)
        async with get_db() as db:
            return await self.plugin_settings_repo.set_config(db, plugin_id, plugin.config)

    def get_plugin(self, plugin_id: str) -> Optional[BaseCrawlerPlugin]:
        return registry.get(plugin_id)

    async def run_plugin(self, plugin_id: str) -> List[ProxyRaw]:
        """执行单个插件爬取"""
        plugin = self.get_plugin(plugin_id)
        if not plugin:
            raise ValueError(f"Plugin {plugin_id} not found")
        if not plugin.enabled:
            logger.warning(f"Plugin {plugin_id} is disabled, skip crawl")
            return []

        try:
            results = await plugin.crawl()
            self._record_stat(plugin_id, success=len(results))
            logger.info(f"Plugin {plugin_id} crawled {len(results)} proxies")
            return results
        except Exception as e:
            self._record_stat(plugin_id, failure=1)
            logger.error(f"Plugin {plugin_id} crawl failed: {e}")
            return []

    async def run_all_plugins(self) -> List[ProxyRaw]:
        """执行所有启用插件的爬取"""
        all_results: List[ProxyRaw] = []
        for plugin in registry.list_plugins():
            if not plugin.enabled:
                continue
            try:
                results = await self.run_plugin(plugin.name)
                all_results.extend(results)
            except Exception as e:
                logger.error(f"Run all plugins error at {plugin.name}: {e}")
        # 去重
        seen = set()
        unique = []
        for p in all_results:
            key = (p.ip, p.port, p.protocol)
            if key not in seen:
                seen.add(key)
                unique.append(p)
        return unique

    def _record_stat(self, plugin_id: str, success: int = 0, failure: int = 0):
        if plugin_id not in self._stats:
            self._stats[plugin_id] = {
                "success_count": 0,
                "failure_count": 0,
                "last_run": None,
            }
        self._stats[plugin_id]["success_count"] += success
        self._stats[plugin_id]["failure_count"] += failure
        if success or failure:
            self._stats[plugin_id]["last_run"] = datetime.now()