import pymysql import sqlite3 import threading import os import json from myutils.ConfigManager import myCongif from myutils.MyLogger_logger import LogHandler from myutils.MyTime import get_local_timestr class DBManager: #实例化数据库管理对象，并连接数据库 #itype=0 使用mysql数据库，1-使用sqlite数据库 def __init__(self): self.logger = LogHandler().get_logger("DBManager") self.lock = threading.Lock() self.itype = myCongif.get_data("DBType") self.ok = False if self.itype ==0: self.host = myCongif.get_data('mysql.host') self.port = myCongif.get_data('mysql.port') self.user = myCongif.get_data('mysql.user') self.passwd = myCongif.get_data('mysql.passwd') self.database = myCongif.get_data('mysql.database') self.connection = None self.cursor = None elif self.itype ==1: self.dbfile = myCongif.get_data("sqlite") if not os.path.exists(self.dbfile): self.dbfile = "../" + self.dbfile #直接运行DBManager时初始路径不是在根目录 if not os.path.exists(self.dbfile): raise FileNotFoundError(f"Database file {self.dbfile} does not exist.") else: self.logger.error("错误的数据库类型，请检查") def __del__(self): if self.ok: self.cursor.close() self.connection.close() self.cursor = None self.connection = None self.logger.debug("DBManager销毁") def connect(self): try: if self.itype ==0: self.connection = pymysql.connect(host=self.host, port=self.port, user=self.user, passwd=self.passwd, db=self.database,charset='utf8') self.cursor = self.connection.cursor() elif self.itype ==1: self.connection = sqlite3.connect(self.dbfile) self.cursor = self.connection.cursor() self.ok = True self.logger.debug("服务器端数据库连接成功") return True except: self.logger.error("服务器端数据库连接失败") return False # 判断数据库连接是否正常，若不正常则重连接 def Retest_conn(self): if self.itype == 0: #除了mysql，sqlite3不需要判断连接状态 try: self.connection.ping() except: return self.connect() return True # 执行数据库查询操作 1-只查询一条记录，其他所有记录 def do_select(self, strsql, itype=0): # self.conn.begin() self.lock.acquire() data = None if self.Retest_conn(): try: self.cursor.execute(strsql) self.connection.commit() # select要commit提交事务，是存在获取不到最新数据的问题（innoDB事务机制） except Exception as e: self.logger.error("do_select异常报错:%s" % str(e)) self.lock.release() return None if itype == 1: data = self.cursor.fetchone() else: data = self.cursor.fetchall() self.lock.release() return data # 执行数据库语句 def do_sql(self, strsql, data=None): bok = False self.lock.acquire() if self.Retest_conn(): try: # self.conn.begin() if data: iret = self.cursor.executemany(strsql, data) #批量执行sql语句 else: iret = self.cursor.execute(strsql) self.connection.commit() bok = True except Exception as e: self.logger.error("执行数据库语句%s出错:%s" % (strsql, str(e))) self.connection.rollback() self.lock.release() return bok def safe_do_sql(self,strsql,params): bok = False self.lock.acquire() if self.Retest_conn(): try: self.cursor.execute(strsql, params) self.connection.commit() bok = True except Exception as e: self.logger.error("执行数据库语句%s出错:%s" % (strsql, str(e))) self.connection.rollback() self.lock.release() return bok def is_json(self,s:str) -> bool: if not isinstance(s, str): return False try: json.loads(s) return True except json.JSONDecodeError: return False except Exception: return False # 处理其他意外异常（如输入 None） #---------------------特定数据库操作函数--------------------- def start_task(self,task_name,task_target) -> int: ''' 数据库添加检测任务 :param task_name: :param task_target: :return: task_id ''' task_id =0 start_time = get_local_timestr() sql = "INSERT INTO task (task_name,task_target,start_time) VALUES (%s,%s,%s)" params = (task_name,task_target,start_time) self.safe_do_sql(sql,params) task_id = self.cursor.lastrowid return task_id #指令执行结果入库 def insetr_result(self,task_id,instruction,result,do_sn,start_time,end_time,source_result,ext_params): # 统一将 result 转为 JSON 字符串（无论原始类型） try: if not isinstance(result, str): str_result = json.dumps(result, ensure_ascii=False) else: # 如果是字符串，先验证是否为合法 JSON（可选） json.loads(result) str_result = result except (TypeError, json.JSONDecodeError): str_result = json.dumps(str(result)) # 兜底处理非 JSON 字符串 # 使用参数化查询 sql = """ INSERT INTO task_result (task_id, instruction, result, do_sn,start_time,end_time,source_result,is_user,is_vulnerability) VALUES (%s, %s, %s, %s, %s, %s,%s,%s,%s) """ params = (task_id, instruction, str_result, do_sn,start_time,end_time,source_result,ext_params['is_user'], ext_params['is_vulnerability']) return self.safe_do_sql(sql,params) #llm数据入库 def insert_llm(self,task_id,prompt,reasoning_content,content,post_time,llm_sn): str_reasoning = "" str_content = "" try: if not isinstance(reasoning_content, str): str_reasoning = json.dumps(reasoning_content, ensure_ascii=False) else: # 如果是字符串，先验证是否为合法 JSON（可选） json.loads(reasoning_content) str_reasoning = reasoning_content except (TypeError, json.JSONDecodeError): str_reasoning = json.dumps(str(reasoning_content)) # 兜底处理非 JSON 字符串 try: if not isinstance(content, str): str_content = json.dumps(content, ensure_ascii=False) else: # 如果是字符串，先验证是否为合法 JSON（可选） json.loads(content) str_content = content except (TypeError, json.JSONDecodeError): str_content = json.dumps(str(content)) # 兜底处理非 JSON 字符串 sql=""" INSERT INTO task_llm (task_id,do_sn,prompt,reasoning_content,content,start_time) VALUES (%s, %s, %s, %s, %s, %s) """ params = (task_id,llm_sn,prompt,str_reasoning,str_content,post_time) return self.safe_do_sql(sql,params) def test(self): # 建立数据库连接 conn = pymysql.connect( host='localhost', port=3306, user='username', password='password', database='database_name' ) # 创建游标对象 cursor = conn.cursor() # 执行 SQL 查询 query = "SELECT * FROM table_name" cursor.execute(query) # 获取查询结果 result = cursor.fetchall() # 输出结果 for row in result: print(row) # 关闭游标和连接 cursor.close() conn.close() if __name__ == "__main__": mDBM = DBManager() mDBM.connect() print(mDBM.start_task("11","22"))