import json import os import logging import sqlite3 from typing import List, Dict, Any, Optional # 导入数据库管理器 from core.database import get_db_manager # 配置日志 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) class RuleManager: ''' 规则管理器，负责规则的增删改查和持久化。 ''' def __init__(self, db_path: str = 'data/honeypot.db', rules_file: str = 'config/rules.json'): self.rules_file = rules_file self.rules: List[Dict[str, Any]] = [] self.db_manager = get_db_manager() # 获取数据库管理器实例 # 重新初始化数据库管理器以使用指定的路径 if db_path != 'data/honeypot.db': from core.database import DatabaseManager self.db_manager = DatabaseManager(db_path) self._load_rules() def _load_rules(self): '''从文件加载全局规则''' if os.path.exists(self.rules_file): try: with open(self.rules_file, 'r', encoding='utf-8') as f: data = json.load(f) # Ensure 'rules' key exists and is a list self.rules = data.get('rules', []) if isinstance(data, dict) else [] logger.info(f"Loaded {len(self.rules)} global rules from {self.rules_file}") except Exception as e: logger.error(f"Failed to load rules from {self.rules_file}: {e}") self.rules = [] else: logger.warning(f"Rules file {self.rules_file} not found. Initializing with empty rules list.") self.rules = [] def _save_rules(self): '''将全局规则保存到文件''' try: # 确保 config 目录存在 os.makedirs(os.path.dirname(self.rules_file), exist_ok=True) with open(self.rules_file, 'w', encoding='utf-8') as f: json.dump({'rules': self.rules}, f, ensure_ascii=False, indent=2) logger.info(f"Saved {len(self.rules)} global rules to {self.rules_file}") except Exception as e: logger.error(f"Failed to save rules to {self.rules_file}: {e}") def get_all_rules(self) -> List[Dict[str, Any]]: '''获取所有全局规则''' return self.rules.copy() def get_rule_by_id(self, rule_id: str) -> Optional[Dict[str, Any]]: '''根据ID获取全局规则''' for rule in self.rules: if rule.get('id') == rule_id: return rule return None def add_rule(self, rule: Dict[str, Any]) -> bool: '''添加新全局规则''' if not rule.get('id'): logger.error("Rule must have an 'id'") return False if self.get_rule_by_id(rule['id']): logger.error(f"Rule with ID '{rule['id']}' already exists") return False self.rules.append(rule) self._save_rules() return True def update_rule(self, rule_id: str, updated_rule: Dict[str, Any]) -> bool: '''更新全局规则''' for i, rule in enumerate(self.rules): if rule.get('id') == rule_id: self.rules[i] = updated_rule self._save_rules() return True logger.warning(f"Rule with ID '{rule_id}' not found for update") return False def delete_rule(self, rule_id: str) -> bool: '''删除全局规则''' for i, rule in enumerate(self.rules): if rule.get('id') == rule_id: del self.rules[i] self._save_rules() return True logger.warning(f"Rule with ID '{rule_id}' not found for deletion") return False def enable_rule(self, rule_id: str) -> bool: '''启用全局规则''' rule = self.get_rule_by_id(rule_id) if rule: rule['enabled'] = True self._save_rules() return True return False def disable_rule(self, rule_id: str) -> bool: '''禁用全局规则''' rule = self.get_rule_by_id(rule_id) if rule: rule['enabled'] = False self._save_rules() return True return False # --- 数据库交互方法 (环境和环境规则) --- def create_environment(self, env_id: str, driver: str, config: Dict[str, Any]) -> bool: """ 在数据库中创建一个新环境。 :param env_id: 环境唯一标识符。 :param driver: 使用的驱动名称。 :param config: 环境配置字典。 :return: 操作是否成功。 """ try: with self.db_manager.get_connection() as conn: cursor = conn.cursor() cursor.execute( "INSERT INTO environments (id, driver, config, status) VALUES (?, ?, ?, ?)", (env_id, driver, json.dumps(config, ensure_ascii=False), "stopped") ) conn.commit() logger.info(f"Environment {env_id} created in database.") return True except sqlite3.IntegrityError: logger.warning(f"Environment {env_id} already exists in database.") return False except sqlite3.Error as e: logger.error(f"Error creating environment {env_id}: {e}") return False def get_environment(self, env_id: str) -> Optional[Dict[str, Any]]: """ 从数据库获取环境信息。 :param env_id: 环境唯一标识符。 :return: 环境信息字典，如果未找到则返回 None。 """ try: with self.db_manager.get_connection() as conn: cursor = conn.cursor() cursor.execute("SELECT * FROM environments WHERE id = ?", (env_id,)) row = cursor.fetchone() if row: # 将 config JSON 字符串转换回字典 env_data = dict(row) env_data['config'] = json.loads(env_data['config']) return env_data except sqlite3.Error as e: logger.error(f"Error getting environment {env_id}: {e}") return None def update_environment(self, env_id: str, status: str) -> bool: try: with self.db_manager.get_connection() as conn: cursor = conn.cursor() cursor.execute( "UPDATE environments set status = ? where id = ?", (status, env_id) ) conn.commit() logger.info(f"Environment {env_id} update in database.") return True except sqlite3.IntegrityError: logger.warning(f"Environment {env_id} already exists in database.") return False except sqlite3.Error as e: logger.error(f"Error updateing environment {env_id}: {e}") return False def list_environments(self) -> List[Dict[str, Any]]: """ 从数据库列出所有环境。 :return: 环境信息字典列表。 """ try: with self.db_manager.get_connection() as conn: cursor = conn.cursor() cursor.execute("SELECT * FROM environments") rows = cursor.fetchall() environments = [] for row in rows: env_data = dict(row) # 将 config JSON 字符串转换回字典 env_data['config'] = json.loads(env_data['config']) environments.append(env_data) return environments except sqlite3.Error as e: logger.error(f"Error listing environments: {e}") return [] def delete_environment(self, env_id: str) -> bool: """ 从数据库删除环境。由于设置了 CASCADE，environment_rules 表中的相关记录也会被自动删除。 :param env_id: 环境唯一标识符。 :return: 操作是否成功。 """ try: with self.db_manager.get_connection() as conn: cursor = conn.cursor() cursor.execute("DELETE FROM environments WHERE id = ?", (env_id,)) conn.commit() if cursor.rowcount > 0: logger.info(f"Environment {env_id} deleted from database.") return True else: logger.warning(f"Environment {env_id} not found for deletion.") return False except sqlite3.Error as e: logger.error(f"Error deleting environment {env_id}: {e}") return False def set_environment_rules(self, env_id: str, rules: List[Dict[str, Any]]) -> bool: """ 为指定环境设置（覆盖）规则列表。 :param env_id: 环境唯一标识符。 :param rules: 要关联到环境的规则对象列表。 :return: 操作是否成功。 """ try: with self.db_manager.get_connection() as conn: cursor = conn.cursor() # 1. 先删除该环境现有的所有规则 cursor.execute("DELETE FROM environment_rules WHERE environment_id = ?", (env_id,)) # 2. 插入新的规则列表 for index, rule in enumerate(rules): rule_json = json.dumps(rule, ensure_ascii=False) cursor.execute( "INSERT INTO environment_rules (environment_id, rule_data, rule_order) VALUES (?, ?, ?)", (env_id, rule_json, index) ) conn.commit() logger.info(f"Set {len(rules)} rules for environment {env_id}.") return True except sqlite3.Error as e: logger.error(f"Error setting rules for environment {env_id}: {e}") return False def get_environment_rules(self, env_id: str) -> List[Dict[str, Any]]: """ 获取指定环境的规则列表。 :param env_id: 环境唯一标识符。 :return: 与该环境关联的规则对象列表。 """ try: with self.db_manager.get_connection() as conn: cursor = conn.cursor() # 按 rule_order 排序以保证顺序 cursor.execute( "SELECT rule_data FROM environment_rules WHERE environment_id = ? ORDER BY rule_order", (env_id,) ) rows = cursor.fetchall() rules = [] for row in rows: # 将 rule_data JSON 字符串转换回字典 rule_obj = json.loads(row['rule_data']) rules.append(rule_obj) return rules except sqlite3.Error as e: logger.error(f"Error getting rules for environment {env_id}: {e}") # 发生错误时返回空列表或全局规则作为后备？ # 这里选择返回空列表，让调用者决定是否使用后备方案 return []