import os import sys from dataclasses import dataclass, field from datetime import datetime from typing import Optional, List, Dict, Tuple, Any, Union, Type, Generic, TypeVar, TextIO import sqlite3 import json if "/www/server/panel/class" not in sys.path: sys.path.insert(0, "/www/server/panel/class") import public import db if "/www/server/panel" not in sys.path: sys.path.insert(0, "/www/server/panel") @dataclass class Script: """对应scripts表""" name: str script_type: str content: str id: Optional[int] = None description: Optional[str] = None group_id: int = 0 created_at: Optional[datetime] = None updated_at: Optional[datetime] = None @staticmethod def check(data: Dict[str, Any]) -> str: if "script_type" not in data or not data["script_type"]: return "Script type cannot be empty" if not data["script_type"] in ["python", "shell"]: return "Script type error, please choose Python or Shell" if "content" not in data or not data["content"]: return "Script content cannot be empty" if "name" not in data or not data["name"]: return "Script name cannot be empty" return "" @classmethod def from_dict(cls, data: Dict[str, Any]) -> 'Script': """从字典创建Script实例""" return cls( id=int(data['id']) if data.get('id', None) else None, name=str(data['name']), script_type=str(data['script_type']), content=str(data['content']), description=str(data['description']) if data.get('description', None) else None, group_id=int(data['group_id']) if data.get('group_id', None) else 0, created_at=datetime.fromisoformat(data['created_at']) if data.get('created_at', None) else None, updated_at=datetime.fromisoformat(data['updated_at']) if data.get('updated_at', None) else None ) def to_dict(self) -> Dict[str, Any]: """转换为字典格式""" return { 'id': self.id, 'name': self.name, 'script_type': self.script_type, 'content': self.content, 'description': self.description, 'group_id': self.group_id, 'created_at': self.created_at.isoformat() if self.created_at else None, 'updated_at': self.updated_at.isoformat() if self.updated_at else None } @dataclass class ScriptGroup: """对应script_groups表""" name: str id: Optional[int] = None description: Optional[str] = None created_at: Optional[datetime] = None @staticmethod def check(data: Dict[str, Any]) -> str: if "name" not in data or not data["name"]: return "Script group name cannot be empty" return "" @classmethod def from_dict(cls, data: Dict[str, Any]) -> 'ScriptGroup': """从字典创建ScriptGroup实例""" return cls( id=int(data['id']) if data.get('id', None) else None, name=str(data['name']), description=str(data['description']) if data.get('description', None) else None, created_at=datetime.fromisoformat(data['created_at']) if data.get('created_at', None) else None ) def to_dict(self) -> Dict[str, Any]: """转换为字典格式""" return { 'id': self.id, 'name': self.name, 'description': self.description, 'created_at': self.created_at.isoformat() if self.created_at else None } @dataclass class ExecutorTask: """对应executor_tasks表""" script_id: int script_content: str script_type: str server_ids: str = "" id: Optional[int] = None created_at: Optional[datetime] = None updated_at: Optional[datetime] = None _elogs: Optional[List["ExecutorLog"]] = None @classmethod def from_dict(cls, data: Dict[str, Any]) -> 'ExecutorTask': """从字典创建ExecutorTask实例""" return cls( id=int(data['id']) if data.get('id', None) else None, script_id=int(data['script_id']), script_content=str(data['script_content']), script_type=str(data['script_type']), created_at=datetime.fromisoformat(data['created_at']) if data.get('created_at', None) else None, updated_at=datetime.fromisoformat(data['updated_at']) if data.get('updated_at', None) else None ) def to_dict(self) -> Dict[str, Any]: """转换为字典格式""" return { 'id': self.id, 'script_id': self.script_id, 'server_ids': self.server_ids, 'script_content': self.script_content, 'script_type': self.script_type, 'created_at': self.created_at.isoformat() if self.created_at else None, 'updated_at': self.updated_at.isoformat() if self.updated_at else None } @property def elogs(self) -> List["ExecutorLog"]: if self._elogs is None: return [] return self._elogs @elogs.setter def elogs(self, elogs: List["ExecutorLog"]): self._elogs = elogs _EXECUTOR_LOG_DIR = public.get_panel_path() + "/logs/executor_log/" try: if not os.path.exists(_EXECUTOR_LOG_DIR): os.makedirs(_EXECUTOR_LOG_DIR) except: pass @dataclass class ExecutorLog: """对应executor_logs表""" executor_task_id: int server_id: int ssh_host: str id: Optional[int] = None status: int = 0 # 0:运行中 1:成功 2:失败 3:异常 log_name: Optional[str] = None created_at: Optional[datetime] = None updated_at: Optional[datetime] = None _log_fp: Optional[TextIO] = None @classmethod def from_dict(cls, data: Dict[str, Any]) -> 'ExecutorLog': """从字典创建ExecutorLog实例""" return cls( id=int(data['id']) if data.get('id', None) else None, executor_task_id=int(data['executor_task_id']), server_id=int(data['server_id']), ssh_host=str(data['ssh_host']), status=int(data['status']) if data.get('status', 0) else 0, log_name=str(data['log_name']) if data.get('log_name', None) else None, created_at=datetime.fromisoformat(data['created_at']) if data.get('created_at', None) else None, updated_at=datetime.fromisoformat(data['updated_at']) if data.get('updated_at', None) else None ) def to_dict(self) -> Dict[str, Any]: """转换为字典格式""" return { 'id': self.id, 'executor_task_id': self.executor_task_id, 'server_id': self.server_id, 'ssh_host': self.ssh_host, 'status': self.status, 'log_name': self.log_name, 'created_at': self.created_at.isoformat() if self.created_at else None, 'updated_at': self.updated_at.isoformat() if self.updated_at else None } @property def log_file(self): return os.path.join(_EXECUTOR_LOG_DIR, self.log_name) @property def log_fp(self): if self._log_fp is None: self._log_fp = open(self.log_file, "w+") return self._log_fp def create_log(self): public.writeFile(self.log_file, "") def remove_log(self): if os.path.exists(self.log_file): os.remove(self.log_file) def get_log(self): return public.readFile(self.log_file) def write_log(self, log_data: str, is_end_log=False): self.log_fp.write(log_data) self.log_fp.flush() if is_end_log: self.log_fp.close() self._log_fp = None _TableType = TypeVar("_TableType", bound=Union[Script, ScriptGroup, ExecutorTask, ExecutorLog]) class _Table(Generic[_TableType]): """数据库表""" table_name: str = "" data_cls: Type[_TableType] def __init__(self, db_obj: db.Sql): self._db = db_obj # 当仅传递一个数据时,返回插入数的 id或错误信息; 当传递多个数据时,返回插入的行数或错误信息 def create(self, data: Union[_TableType, List[_TableType]]) -> Union[int, str]: """创建数据""" if not isinstance(data, list): data = [data] if not len(data): raise ValueError("Data cannot be empty") if not isinstance(data[0], self.data_cls): raise ValueError("Data type error") now = datetime.now().isoformat() def fileter_data(item): item_dict = item.to_dict() if "id" in item_dict: item_dict.pop("id") if "created_at" in item_dict and item_dict["created_at"] is None: item_dict["created_at"] = now if "updated_at" in item_dict and item_dict["updated_at"] is None: item_dict["updated_at"] = now return item_dict data_list = list(map(fileter_data, data)) if len(data_list) == 1: try: res = self._db.table(self.table_name).insert(data_list[0]) if isinstance(res, int): return res return str(res) except Exception as e: return str(e) try: res = self._db.table(self.table_name).batch_insert(data_list) if isinstance(res, (int, bool)): return len(data) return str(res) except Exception as e: return str(e) def update(self, data: _TableType) -> str: """更新数据""" if not isinstance(data, self.data_cls): raise ValueError("Data type error") data_dict = data.to_dict() data_dict.pop('created_at', None) if "updated_at" in data_dict: data_dict["updated_at"] = datetime.now().isoformat() if "id" not in data_dict: raise ValueError("The data ID cannot be empty") try: self._db.table(self.table_name).where("id=?", (data_dict["id"],)).update(data_dict) except Exception as e: return str(e) return "" def get_byid(self, data_id: int) -> Optional[_TableType]: """根据id获取数据""" try: result = self._db.table(self.table_name).where("id=?", (data_id,)).find() except Exception as e: return None if not result: return None return self.data_cls.from_dict(result) def delete(self, data_id: Union[int, List[int]]): """删除数据""" if isinstance(data_id, list): data_id = [int(item) for item in data_id] elif isinstance(data_id, int): data_id = [int(data_id)] else: return "数据id类型错误" try: self._db.table(self.table_name).where( "id in ({})".format(",".join(["?"] * len(data_id))), (*data_id,) ).delete() return "" except Exception as e: return str(e) def query(self, *args) -> List[_TableType]: """查询数据""" try: result = self._db.table(self.table_name).where(*args).select() except Exception as e: return [] if not result: return [] return [self.data_cls.from_dict(item) for item in result] def query_page(self, *args, page_num: int = 1, limit: int = 10) -> List[_TableType]: """查询数据, 支持分页""" try: offset = limit * (page_num - 1) result = self._db.table(self.table_name).where(*args).limit(limit, offset).order("id DESC").select() except Exception as e: public.print_error() return [] if not result: return [] return [self.data_cls.from_dict(item) for item in result] def count(self, *args) -> int: """查询数据数量""" try: result = self._db.table(self.table_name).where(*args).count() except Exception as e: return 0 return result def find(self, *args) -> Optional[_TableType]: """查询单条数据""" try: result = self._db.table(self.table_name).where(*args).find() except Exception as e: return None if not result: return None return self.data_cls.from_dict(result) class _ScriptTable(_Table[Script]): """脚本表""" table_name = "scripts" data_cls = Script def set_group_id(self, group_id: int, *where_args) -> str: """设置脚本组""" try: self._db.table(self.table_name).where(where_args).update({"group_id": group_id}) except Exception as e: return str(e) return "" class _ScriptGroupTable(_Table[ScriptGroup]): """脚本组表""" table_name = "script_groups" data_cls = ScriptGroup default_group = ScriptGroup( id=0, name="default", description="Default grouping, use this grouping when not set", created_at=datetime.now(), ) def all_group(self) -> List[ScriptGroup]: """获取所有脚本组""" try: result = self._db.table(self.table_name).select() except Exception as e: return [] if not result: return [] return [self.default_group] + [self.data_cls.from_dict(item) for item in result] class _ExecutorTaskTable(_Table[ExecutorTask]): """执行任务表""" table_name = "executor_tasks" data_cls = ExecutorTask def query_tasks(self, page=1, size=10, node_id: int = None, script_type: str = None, search: str = None ) -> Tuple[int, List[ExecutorTask]]: """查询任务""" where_args, parms = [], [] if script_type and script_type != "all": where_args.append("script_type=?") parms.append(script_type) if search: search_str = "script_content like ?" parms.append("%{}%".format(search)) stable = _ScriptTable(self._db) data = stable.query("name like ? or description like ?", ("%{}%".format(search), "%{}%".format(search))) if data: search_str += " or script_id in ({})".format(",".join(["?"] * len(data))) where_args.append("(" + search_str + ")") parms.append(tuple([item.id for item in data])) else: where_args.append(search_str) if node_id: where_args.append("server_ids like ?") parms.append("%|{}%".format(node_id)) # public.print_log("search criteria: {}".format(" AND ".join(where_args)), parms) count = self.count( " AND ".join(where_args), (*parms, ) ) return count, self.query_page( " AND ".join(where_args), (*parms, ), page_num=page, limit=size ) class _ExecutorLogTable(_Table[ExecutorLog]): """执行日志表""" table_name = "executor_logs" data_cls = ExecutorLog class ExecutorDB: _DB_FILE = public.get_panel_path() + "/data/db/executor.db" _DB_INIT_FILE = os.path.dirname(__file__) + "/executor.sql" def __init__(self): sql = db.Sql() sql._Sql__DB_FILE = self._DB_FILE self.db = sql self.Script = _ScriptTable(self.db) self.ScriptGroup = _ScriptGroupTable(self.db) self.ExecutorTask = _ExecutorTaskTable(self.db) self.ExecutorLog = _ExecutorLogTable(self.db) def init_db(self): sql_data = public.readFile(self._DB_INIT_FILE) if not os.path.exists(self._DB_FILE) or os.path.getsize(self._DB_FILE) == 0: public.writeFile(self._DB_FILE, "") import sqlite3 conn = sqlite3.connect(self._DB_FILE) cursor = conn.cursor() cursor.executescript(sql_data) conn.commit() conn.close() def close(self): self.db.close() def __enter__(self): return self def __exit__(self, exc_type, exc_value, exc_trackback): self.close()