# -*- coding: utf-8 -*- import importlib import threading import time import uuid from typing import Callable, Dict, Optional from apscheduler.schedulers.background import BackgroundScheduler from Lib.baseplaybook import BasePlaybook from Lib.log import logger from Lib.moduleengine import ModuleEngine from Lib.playbookloader import PlaybookLoader from Lib.threadmodulemanager import thread_module_manager from Lib.xcache import Xcache from PLUGINS.Redis.redis_stream_api import RedisStreamAPI from PLUGINS.SIRP.analysis import run_case_analysis from PLUGINS.SIRP.sirpapi import Playbook, Case from PLUGINS.SIRP.sirpextramodel import PlaybookJobStatus, PlaybookModel class MainMonitor(object): @staticmethod def on_playbook_task_finished(thread_id: str, task_obj: BasePlaybook, result: object, exception: Optional[Exception], context: Optional[Dict[str, str]] = None): context = context or {} playbook_row_id = context.get("playbook_row_id") if not isinstance(playbook_row_id, str) or playbook_row_id == "": logger.error(f"[Thread {thread_id}] Missing playbook_row_id in callback context.") return try: playbook_current = Playbook.get(playbook_row_id, lazy_load=True) except Exception as e: logger.error(f"[Thread {thread_id}] Failed to load playbook for fallback status update.") logger.exception(e) return # Business code has already written a terminal state; do not overwrite. if playbook_current.job_status != PlaybookJobStatus.RUNNING: logger.info( f"[Thread {thread_id}] Skip fallback status update, current status: {playbook_current.job_status}, " f"row_id: {playbook_row_id}" ) return fallback_status = PlaybookJobStatus.SUCCESS if exception is None else PlaybookJobStatus.FAILED if exception is None: remark = f"source=thread_fallback; thread_id={thread_id}; message=Task finished without explicit business status update." else: remark = ( f"source=thread_fallback; thread_id={thread_id}; " f"exception={type(exception).__name__}; message={exception}" ) model_tmp = PlaybookModel(row_id=playbook_row_id) model_tmp.job_status = fallback_status model_tmp.remark = remark Playbook.update(model_tmp) logger.info( f"[Thread {thread_id}] Applied fallback status update: {fallback_status}, row_id: {playbook_row_id}" ) MainScheduler: BackgroundScheduler def __init__(self): self._background_threads: Dict[str, list] = {} self.engine = ModuleEngine() self.redis_stream_api = RedisStreamAPI() self.MainScheduler = BackgroundScheduler(timezone='Asia/Shanghai') @staticmethod def run_task_in_loop(task_func: Callable, task_name: str, retry_interval: int = 5, exec_interval: int = None): """ Run a task function in an infinite loop with error handling Args: task_func: The function to run task_name: Name of the task for logging retry_interval: Seconds to wait between retries on error exec_interval: Seconds to wait between executions (defaults to retry_interval if None) """ # If exec_interval is not specified, use retry_interval if exec_interval is None: exec_interval = retry_interval while True: try: task_func() # Wait for the specified execution interval before running again time.sleep(exec_interval) except Exception as e: logger.error(f"Error in {task_name}") logger.exception(e) time.sleep(retry_interval) def start_background_task(self, task_func: Callable, task_name: str, retry_interval: int = 5, exec_interval: int = None, thread_count: int = 1): """ Start a background task in one or more separate threads Args: task_func: The function to run task_name: Name of the task for logging retry_interval: Seconds to wait between retries on error exec_interval: Seconds to wait between executions (defaults to retry_interval if None) thread_count: Number of threads to spawn for this task """ if task_name not in self._background_threads: self._background_threads[task_name] = [] for i in range(thread_count): suffix = f"#{i}" if thread_count > 1 else "" thread_name = f"{task_name}{suffix}" thread = threading.Thread( target=self.run_task_in_loop, args=(task_func, thread_name, retry_interval, exec_interval), daemon=True, name=thread_name, ) self._background_threads[task_name].append(thread) thread.start() logger.info(f"Started background task: {thread_name}") def start(self): logger.info("Starting background services...") logger.info("Load PlaybookLoader module config") PlaybookLoader.load_all_playbook_config() delay_time = 3 # Start background tasks self.start_background_task(self.subscribe_pending_playbook, "subscribe_pending_playbook", delay_time) self.start_background_task(self.subscribe_case_analysis_scheduler, "subscribe_case_analysis_scheduler", delay_time) self.start_background_task(self.subscribe_case_analysis_queue, "subscribe_case_analysis_queue", delay_time, thread_count=3) # engine self.engine.start() logger.info("Background services started.") @staticmethod def subscribe_pending_playbook(): models = Playbook.list_pending_playbooks() for model in models: module_config = Xcache.get_module_config_by_name(model.name) model_tmp = PlaybookModel(row_id=model.row_id) if module_config is None: PlaybookLoader.load_all_playbook_config() # try again module_config = Xcache.get_module_config_by_name(model.name) if module_config is None: logger.error(f"PlaybookLoader module config not found: {model.name}") model_tmp.job_status = PlaybookJobStatus.FAILED model_tmp.remark = f"PlaybookLoader module config not found: {model.name}" Playbook.update(model_tmp) continue load_path = module_config.get("load_path") try: class_intent = importlib.import_module(load_path) playbook_intent: BasePlaybook = class_intent.Playbook() playbook_intent._playbook_model = model except Exception as E: logger.exception(E) model_tmp.job_status = PlaybookJobStatus.FAILED model_tmp.remark = str(E) Playbook.update(model_tmp) continue job_id = str(uuid.uuid1()) model_tmp.job_status = PlaybookJobStatus.RUNNING model_tmp.job_id = job_id Playbook.update(model_tmp) try: thread_module_manager.start_task( playbook_intent, thread_id=job_id, on_finished=MainMonitor.on_playbook_task_finished, callback_context={ "playbook_row_id": model.row_id, }, ) except Exception as e: model_tmp.job_status = PlaybookJobStatus.FAILED model_tmp.remark = f"Failed to create playbook job. exception={type(e).__name__}; message={e}" Playbook.update(model_tmp) continue logger.info(f"Create playbook job success: {job_id}") @staticmethod def subscribe_case_analysis_scheduler(): # Periodically move due cases into the analysis queue. # 周期性扫描到点案件，并将其送入分析队列。 promoted_row_ids = Case.promote_due_analysis_cases() if promoted_row_ids: logger.info(f"Queued {len(promoted_row_ids)} case(s) for scheduled analysis.") def subscribe_case_analysis_queue(self): # Consume one queue message and hand it to the analysis runner. # 每次消费一条队列消息，并交给分析执行器处理。 message = self.redis_stream_api.read_stream_message_with_id( stream_name=Case.ANALYSIS_STREAM_NAME, consumer_name="case-analysis-worker", ) queue_message_id = message.get("message_id") payload = message.get("data", {}) case_row_id = payload.get("case_row_id") trigger = payload.get("trigger", "unknown") if not case_row_id: logger.error("Case analysis queue message missing case_row_id.") return run_case_analysis(case_row_id, trigger, queue_message_id=queue_message_id)