You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
908 lines
34 KiB
908 lines
34 KiB
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
|
|
from datetime import timedelta
|
|
from datetime import datetime, timedelta
|
|
|
|
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.COLUMN_LIMITS = {} #各个表字段长度信息
|
|
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
|
|
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.connection.close()
|
|
self.connection = None
|
|
self.logger.debug("DBManager销毁")
|
|
|
|
def get_column_limits(self):
|
|
'''连接成功后，统一读取数据库字段列长度设定'''
|
|
tables = ['assets_user','ip_assets','ip_to_url','port_assets','task','task_llm','task_result','task_vul','url_assets','user','zf_system']
|
|
for table in tables:
|
|
table_limit = {}
|
|
strsql = f'''
|
|
SELECT COLUMN_NAME, CHARACTER_MAXIMUM_LENGTH
|
|
FROM information_schema.columns
|
|
WHERE table_schema = 'your_db'
|
|
AND table_name = '{table}';
|
|
'''
|
|
datas = self.do_select(strsql)
|
|
for data in datas:
|
|
if data[1]:
|
|
table_limit[f"{str(data[0])}"] = data[1]
|
|
self.COLUMN_LIMITS[f"{table}"] = table_limit
|
|
|
|
def trim_fields(self,table: str, data: dict, logger) -> dict:
|
|
"""
|
|
按 COLUMN_LIMITS 截断过长的 varchar/text 字段。
|
|
:param table: 表名
|
|
:param data : 待插入/更新的字段 dict
|
|
:return : 新 dict（已截断）
|
|
"""
|
|
limits = self.COLUMN_LIMITS.get(table, {})
|
|
trimmed = {}
|
|
|
|
for col, val in data.items():
|
|
if val is None or col not in limits:
|
|
trimmed[col] = val
|
|
continue
|
|
|
|
max_len = limits[col]
|
|
# 只处理 str 类型
|
|
if isinstance(val, str) and len(val) > max_len:
|
|
self.logger.debug(
|
|
f"{table}.{col} 超长：{len(val)} > {max_len}，已截断"
|
|
)
|
|
trimmed[col] = val[:max_len]
|
|
else:
|
|
trimmed[col] = val
|
|
return trimmed
|
|
|
|
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')
|
|
elif self.itype ==1:
|
|
self.connection = sqlite3.connect(self.dbfile)
|
|
self.ok = True
|
|
self.logger.debug("服务器端数据库连接成功")
|
|
#self.get_column_limits()
|
|
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.connection.commit() # select要commit提交事务，是存在获取不到最新数据的问题（innoDB事务机制）
|
|
with self.connection.cursor() as cursor:
|
|
cursor.execute(strsql)
|
|
if itype == 1:
|
|
data = cursor.fetchone()
|
|
else:
|
|
data = cursor.fetchall()
|
|
except Exception as e:
|
|
self.logger.error("do_select异常报错:%s" % str(e))
|
|
self.lock.release()
|
|
return None
|
|
self.lock.release()
|
|
return data
|
|
|
|
# 执行数据库语句
|
|
def do_sql(self, strsql, data=None):
|
|
bok = False
|
|
self.lock.acquire()
|
|
if self.Retest_conn():
|
|
try:
|
|
with self.connection.cursor() as cursor:
|
|
# self.conn.begin()
|
|
if data:
|
|
iret = cursor.executemany(strsql, data) #批量执行sql语句
|
|
else:
|
|
iret = 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,itype=0,table=None,field_names=None):
|
|
"""
|
|
table : 目标表名（若要自动截断必须填）
|
|
field_names : 与 params 顺序对应的列名列表
|
|
"""
|
|
if table and field_names:
|
|
data_map = dict(zip(field_names, params))
|
|
data_map = self.trim_fields(table, data_map, self.logger)
|
|
params = tuple(data_map[col] for col in field_names)
|
|
|
|
bok = False
|
|
do_id = 0
|
|
self.lock.acquire()
|
|
if self.Retest_conn():
|
|
try:
|
|
with self.connection.cursor() as cursor:
|
|
cursor.execute(strsql, params)
|
|
self.connection.commit()
|
|
if itype ==1: #取insert 的自增id
|
|
do_id = cursor.lastrowid
|
|
elif itype ==2: #取删除的id,需要添加RETURNING id;
|
|
row = cursor.fetchone()
|
|
if row:
|
|
do_id = row[0]
|
|
bok = True
|
|
except Exception as e:
|
|
self.logger.error("执行数据库语句%s出错:%s" % (strsql, str(e)))
|
|
self.connection.rollback()
|
|
self.lock.release()
|
|
return bok,do_id
|
|
|
|
def safe_do_select(self,strsql,params,itype=0):
|
|
results = []
|
|
self.lock.acquire()
|
|
if self.Retest_conn():
|
|
self.connection.commit()
|
|
try:
|
|
with self.connection.cursor() as cursor:
|
|
cursor.execute(strsql, params) # 执行参数化查询
|
|
if itype ==0:
|
|
results = cursor.fetchall() # 获取所有结果
|
|
elif itype ==1:
|
|
results = cursor.fetchone() #获得一条记录
|
|
except Exception as e:
|
|
print(f"查询出错: {e}--\n{strsql}")
|
|
self.lock.release()
|
|
return results
|
|
|
|
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 timedelta_to_str(delta: timedelta) -> str:
|
|
hours, remainder = divmod(delta.total_seconds(), 3600)
|
|
minutes, seconds = divmod(remainder, 60)
|
|
return f"{int(hours):02}:{int(minutes):02}:{int(seconds):02}"
|
|
|
|
#---------------------特定数据库操作函数---------------------
|
|
def get_system_info(self):
|
|
strsql = "select * from zf_system;"
|
|
data = self.do_select(strsql,1)
|
|
return data
|
|
|
|
def get_run_tasks(self):
|
|
strsql = "select ID,task_target,task_status,work_type,cookie_info,llm_type,safe_rank,fake_target from task where task_status <> 2 order by ID;"
|
|
datas = self.do_select(strsql)
|
|
return datas
|
|
|
|
def start_task(self,test_target,cookie_info,work_type,llm_type,fake_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,task_status,safe_rank,work_type,cookie_info,llm_type,fake_target) " \
|
|
"VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s)"
|
|
params = (test_target,test_target,start_time,1,0,work_type,cookie_info,llm_type,fake_target)
|
|
bok,task_id = self.safe_do_sql(sql,params,1)
|
|
return task_id
|
|
|
|
def over_task(self,task_id):
|
|
over_time = get_local_timestr()
|
|
strsql = "update task set task_status=2,end_time=%s where ID=%s;"
|
|
params = (over_time,task_id)
|
|
bok,_ = self.safe_do_sql(strsql, params)
|
|
return bok
|
|
|
|
#0<-->1
|
|
def update_task_status(self,task_id,new_status):
|
|
strsql = "update task set task_status=%s where ID=%s;"
|
|
params = (new_status,task_id)
|
|
bok, _ = self.safe_do_sql(strsql, params)
|
|
return bok
|
|
|
|
def update_task_work_type(self,task_id,new_work_type):
|
|
strsql = "update task set work_type=%s where ID=%s;"
|
|
params = (new_work_type, task_id)
|
|
bok, _ = self.safe_do_sql(strsql, params)
|
|
return bok
|
|
|
|
def del_task(self,task_id):
|
|
params = (task_id)
|
|
strsql = "delete from task where ID=%s;"
|
|
bok,_ = self.safe_do_sql(strsql,params)
|
|
strsql = "delete from task_llm where task_id=%s;"
|
|
bok, _ = self.safe_do_sql(strsql, params)
|
|
strsql = "delete from task_result where task_id=%s;"
|
|
bok, _ = self.safe_do_sql(strsql, params)
|
|
strsql = "delete from task_vul where task_id=%s;"
|
|
bok, _ = self.safe_do_sql(strsql, params)
|
|
return bok
|
|
|
|
#指令执行结果入库
|
|
def insetr_result(self,task_id,instruction,result,do_sn,start_time,end_time,source_result,ext_params,node_path):
|
|
str_result = ""
|
|
str_source_result = ""
|
|
# 统一将 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),ensure_ascii=False) # 兜底处理非 JSON 字符串
|
|
|
|
try:
|
|
if not isinstance(source_result, str):
|
|
str_source_result = json.dumps(source_result, ensure_ascii=False)
|
|
else:
|
|
# 如果是字符串，先验证是否为合法 JSON（可选）
|
|
json.loads(source_result)
|
|
str_source_result = source_result
|
|
except (TypeError, json.JSONDecodeError):
|
|
str_source_result = json.dumps(str(source_result),ensure_ascii=False) # 兜底处理非 JSON 字符串
|
|
|
|
# 使用参数化查询
|
|
sql = """
|
|
INSERT INTO task_result
|
|
(task_id, instruction, result, do_sn,start_time,end_time,source_result,is_user,is_vulnerability,node_path)
|
|
VALUES
|
|
(%s, %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'],node_path)
|
|
bok,_ = self.safe_do_sql(sql,params)
|
|
return bok
|
|
|
|
#llm数据入库
|
|
def insert_llm(self,task_id,prompt,reasoning_content,content,post_time,llm_sn,path):
|
|
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)
|
|
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,node_path)
|
|
VALUES
|
|
(%s, %s, %s, %s, %s, %s,%s)
|
|
"""
|
|
str_reasoning = str_reasoning.encode('utf-8').decode('unicode_escape')
|
|
str_content = str_content.encode('utf-8').decode('unicode_escape')
|
|
|
|
params = (task_id,llm_sn,prompt,str_reasoning,str_content,post_time,path)
|
|
bok,_=self.safe_do_sql(sql,params)
|
|
return bok
|
|
|
|
#获取任务的测试指令执行情况
|
|
def get_task_instrs(self,task_id,nodename):
|
|
strsql = '''
|
|
select ID,node_path,do_sn,instruction,result from task_result where task_id = %s
|
|
'''
|
|
params = [task_id]
|
|
if nodename.strip():
|
|
strsql += " and node_path like %s"
|
|
params.append(f"%{nodename}%") # 在参数中添加通配符
|
|
datas = self.safe_do_select(strsql,tuple(params))
|
|
return datas
|
|
|
|
#插入漏洞数据
|
|
def insert_taks_vul(self,task_id,node_name,node_path,vul_type,vul_level,vul_info):
|
|
strsql = '''
|
|
INSERT INTO task_vul
|
|
(task_id,node_name,node_path,vul_type,vul_level,vul_info)
|
|
VALUES (%s,%s,%s,%s,%s,%s)
|
|
'''
|
|
params = (task_id,node_name,node_path,vul_type,vul_level,vul_info)
|
|
bok,_ = self.safe_do_sql(strsql,params)
|
|
return bok
|
|
|
|
#获取任务的漏洞检测情况
|
|
def get_task_vul(self,task_id,nodename,vultype,vullevel):
|
|
strsql = '''
|
|
select ID,node_path,vul_type,vul_level,vul_info from task_vul
|
|
'''
|
|
# 动态构建查询条件
|
|
conditions = ["task_id=%s"] # task_id 必须存在
|
|
params = [task_id] # 参数列表初始化
|
|
|
|
# 按需添加其他条件
|
|
if nodename and nodename.strip(): # 检查nodename是否非空（去除前后空格后）
|
|
conditions.append("node_path like %s")
|
|
params.append(f"%{nodename}%")
|
|
|
|
if vultype and vultype.strip(): # 检查vultype是否非空
|
|
conditions.append("vul_type like %s")
|
|
params.append(f"%{vultype}%")
|
|
|
|
if vullevel and vullevel.strip(): # 检查vullevel是否非空
|
|
conditions.append("vul_level=%s")
|
|
params.append(vullevel)
|
|
|
|
# 组合完整的WHERE子句
|
|
if len(conditions) > 0:
|
|
strsql += " WHERE " + " AND ".join(conditions)
|
|
|
|
# 执行查询（将参数转为元组）
|
|
datas = self.safe_do_select(strsql, tuple(params))
|
|
return datas
|
|
|
|
#获取该任务该节点的所有 已经执行的任务
|
|
def get_task_node_done_instr(self,task_id,nodepath):
|
|
strsql = '''
|
|
select instruction,start_time,result from task_result where task_id=%s and node_path=%s order by start_time desc;
|
|
'''
|
|
params = (task_id,nodepath)
|
|
datas = self.safe_do_select(strsql,params)
|
|
return datas
|
|
|
|
def get_his_tasks(self,target_name,safe_rank,llm_type,start_time,end_time):
|
|
strsql = "select ID,task_target,safe_rank,llm_type,start_time,end_time from task"
|
|
conditions = ["task_status=%s"]
|
|
params = [2]
|
|
# 按需添加其他条件
|
|
if target_name and target_name.strip(): # 检查nodename是否非空（去除前后空格后）
|
|
conditions.append("task_target=%s")
|
|
params.append(target_name)
|
|
|
|
if safe_rank and safe_rank.strip(): # 检查vultype是否非空
|
|
conditions.append("safe_rank=%s")
|
|
params.append(safe_rank)
|
|
|
|
if llm_type and llm_type.strip(): # 检查vullevel是否非空
|
|
conditions.append("llm_type=%s")
|
|
params.append(llm_type)
|
|
|
|
if start_time and start_time.strip(): # 检查vultype是否非空
|
|
conditions.append("start_time >= %s")
|
|
start_date = datetime.strptime(start_time, "%Y-%m-%d")
|
|
# 生成起始时间字符串（当日 00:00:00）
|
|
start_time_str = start_date.strftime("%Y-%m-%d 00:00:00")
|
|
params.append(start_time_str)
|
|
|
|
if end_time and end_time.strip(): # 检查vullevel是否非空
|
|
conditions.append("start_time < %s")
|
|
# 将输入字符串转为日期对象
|
|
end_date = datetime.strptime(end_time, "%Y-%m-%d")
|
|
# 生成结束时间字符串（次日 00:00:00）
|
|
end_time_str = (end_date + timedelta(days=1)).strftime("%Y-%m-%d 00:00:00")
|
|
params.append(end_time_str)
|
|
|
|
# 组合完整的WHERE子句
|
|
if len(conditions) > 0:
|
|
strsql += " WHERE " + " AND ".join(conditions)
|
|
strsql += " order by start_time DESC"
|
|
|
|
# 执行查询（将参数转为元组）
|
|
datas = self.safe_do_select(strsql, tuple(params))
|
|
return datas
|
|
|
|
def getsystem_info(self):
|
|
strsql = "select local_ip,version from zf_system;"
|
|
data = self.do_select(strsql,1)
|
|
return data
|
|
|
|
def update_localip(self,local_ip):
|
|
strsql = "update zf_system set local_ip=%s;"
|
|
params = (local_ip)
|
|
bok,_ = self.safe_do_sql(strsql,params)
|
|
return bok
|
|
|
|
def get_one_instr(self,instr_id):
|
|
strsql = "select instruction from task_result where ID = %s"
|
|
params = (instr_id)
|
|
data = self.safe_do_select(strsql,params,1)
|
|
return data[0]
|
|
|
|
#------------------资产表相关----------------------
|
|
def add_or_update_URL_asset(self,Domain,Subdomains,Registrant,Email,Creation_date,Expiration_date,Registrar):
|
|
strsql = "select ID from url_assets where URL= %s"
|
|
params = (Domain)
|
|
data = self.safe_do_select(strsql,params,1)
|
|
do_time = get_local_timestr()
|
|
if Subdomains:
|
|
strSubdomains = ','.join(Subdomains)
|
|
else:
|
|
strSubdomains = ""
|
|
if not data:#没有数据则新增
|
|
strsql = "insert into url_assets (registrar,creation_date,expiration_date,emails,create_time,update_time,URL,subdomains,Registrant) " \
|
|
"values (%s,%s,%s,%s,%s,%s,%s,%s,%s)"
|
|
params =(Registrar,Creation_date,Expiration_date,Email,do_time,do_time,Domain,strSubdomains,Registrant)
|
|
bok,url_id = self.safe_do_sql(strsql,params,1)
|
|
else:#有值则修改
|
|
url_id = data[0]
|
|
strsql = "update url_assets set registrar=%s,Registrant=%s,creation_date=%s,expiration_date=%s,emails=%s,update_time=%s," \
|
|
"subdomains=%s where ID = %s;"
|
|
params = (Registrar,Registrant,Creation_date,Expiration_date,Email,do_time,strSubdomains,url_id)
|
|
bok,_ = self.safe_do_sql(strsql,params)
|
|
#维护历史记录 --由数据库触发器维护
|
|
return url_id
|
|
|
|
def add_or_update_IP_asset(self,IP,ip_type):
|
|
#return ip_id scan_count
|
|
strsql = "select id,scan_count from ip_assets where ip_address = %s;"
|
|
params = (IP)
|
|
data = self.safe_do_select(strsql,params,1)
|
|
if data:
|
|
return data[0],data[1]
|
|
#IP没有则新建入库
|
|
ip_id = 0
|
|
start_time = get_local_timestr()
|
|
sql = "INSERT INTO ip_assets (ip_address,ip_version,created_time,scan_count) VALUES (%s,%s,%s,%s);"
|
|
params = (IP, ip_type, start_time,0)
|
|
bok, ip_id = self.safe_do_sql(sql, params, 1)
|
|
return ip_id,0
|
|
|
|
#将task_id 和ip资产进行关联
|
|
def add_task_to_ip(self,task_id,ip_id):
|
|
strsql = "INSERT INTO task_to_ip (task_id,ip_id) VALUES (%s,%s);"
|
|
params = (task_id,ip_id)
|
|
bok,_ = self.safe_do_sql(strsql,params)
|
|
return bok
|
|
|
|
def update_port(self,ip_id,scan_count,Prots):
|
|
update_time = get_local_timestr()
|
|
scan_count += 1
|
|
strsql = "update ip_assets set update_time = %s,scan_count = %s where id=%s"
|
|
params = (update_time,scan_count,ip_id)
|
|
bok,_ = self.safe_do_sql(strsql,params)
|
|
if bok:
|
|
#最新的port数据入库 {\"Protocol\":\"TCP/UDP\",\"Status\":\"open/closed/filtered\"}；
|
|
for port in Prots:
|
|
p_num = port["Port"]
|
|
service = port["Service"]
|
|
version = port["Version"]
|
|
protocol = port["Protocol"]
|
|
status = port["Status"]
|
|
strsql = "insert into port_assets (port,service,version,status,ip_id,scan_count,scan_time,Protocol) " \
|
|
"values (%s,%s,%s,%s,%s,%s,%s,%s)"
|
|
params = (p_num,service,version,status,ip_id,scan_count,update_time,protocol)
|
|
bok,_ = self.safe_do_sql(strsql,params)
|
|
return bok
|
|
|
|
def update_url_to_ip(self,url_id,ips):
|
|
strsql = "select ip_id from ip_to_url where url_id = %s"
|
|
params = (url_id)
|
|
datas = self.safe_do_select(strsql,params)
|
|
old_ips = []
|
|
for data in datas:
|
|
old_ips.append(data[0])
|
|
only_in_old = list(set(old_ips) - set(ips)) #适合不重复，不关心顺序的情况
|
|
only_in_new = list(set(ips) - set(old_ips))
|
|
do_time = get_local_timestr()
|
|
if only_in_old:#新的关联中没有，老的有，就是删除了
|
|
placeholders = ",".join(["(%s,%s,%s)"] * len(only_in_old))
|
|
sql = f"""
|
|
INSERT INTO ip_to_url_his (ip_id, url_id, del_time)
|
|
VALUES {placeholders}
|
|
"""
|
|
# 扁平化参数列表
|
|
params: list = []
|
|
for ip in only_in_old:
|
|
params += [ip, url_id, do_time]
|
|
# 一次性执行
|
|
bok, _ = self.safe_do_sql(sql, tuple(params))
|
|
if not bok:
|
|
raise RuntimeError("批量插入 ip_to_url_his 失败")
|
|
#老表中删除记录---待验证
|
|
strsql = '''
|
|
delete from ip_to_url where url_id=%s and ip_id in (%s)
|
|
'''
|
|
del_ips = ','.join([str(x) for x in only_in_old])
|
|
params = (url_id,del_ips)
|
|
bok,_ = self.safe_do_sql(strsql,params)
|
|
if not bok:
|
|
raise RuntimeError("批量删除 ip_to_url 失败")
|
|
|
|
if only_in_new:#新的有，老的没有就是新增。
|
|
placeholders = ",".join(["(%s,%s,%s)"] * len(only_in_new))
|
|
sql = f"""
|
|
INSERT INTO ip_to_url (ip_id, url_id, create_time)
|
|
VALUES {placeholders}
|
|
"""
|
|
# 扁平化参数列表
|
|
params: list = []
|
|
for ip in only_in_new:
|
|
params += [ip, url_id, do_time]
|
|
# 一次性执行
|
|
bok, _ = self.safe_do_sql(sql, tuple(params))
|
|
if not bok:
|
|
raise RuntimeError("批量插入 ip_to_url 失败")
|
|
|
|
def get_ip_assets_db(self,IP,user,risk_rank):
|
|
strsql = '''
|
|
SELECT
|
|
ia.ip_address,
|
|
au.uname,
|
|
ia.risk_rank,
|
|
ia.update_time,
|
|
COALESCE(p.port_cnt ,0) AS port_total,
|
|
COALESCE(u.url_cnt ,0) AS url_total
|
|
FROM ip_assets AS ia
|
|
LEFT JOIN assets_user AS au ON ia.owner_id = au.ID
|
|
/* 端口数量 */
|
|
LEFT JOIN (
|
|
SELECT ip_id,scan_count,COUNT(*) AS port_cnt
|
|
FROM port_assets
|
|
GROUP BY ip_id,scan_count
|
|
) AS p ON p.ip_id = ia.id and p.scan_count=ia.scan_count
|
|
/* URL 数量 */
|
|
LEFT JOIN (
|
|
SELECT ip_id, COUNT(*) AS url_cnt
|
|
FROM ip_to_url
|
|
GROUP BY ip_id
|
|
) AS u ON u.ip_id = ia.id
|
|
WHERE
|
|
(%s IS NULL OR ia.ip_address LIKE %s)
|
|
AND (%s IS NULL OR au.uname LIKE %s)
|
|
AND (%s IS NULL OR ia.risk_rank = %s);
|
|
'''
|
|
# 构造参数
|
|
ip_like = f"%{IP}%" if IP else None
|
|
user_like = f"%{user}%" if user else None
|
|
rk = risk_rank
|
|
|
|
params = (
|
|
ip_like, ip_like,
|
|
user_like, user_like,
|
|
rk, rk,
|
|
)
|
|
# cursor.execute(sql, params)
|
|
# rows = cursor.fetchall()
|
|
datas = self.safe_do_select(strsql,params)
|
|
return datas
|
|
|
|
def get_ip_info_db(self,IP):
|
|
strsql = '''
|
|
select au.ID,au.tellnum,au.tell_username from assets_user as au
|
|
left join ip_assets as ai on ai.owner_id = au.ID
|
|
where ai.ip_address = %s;
|
|
'''
|
|
params = (IP)
|
|
data = self.safe_do_select(strsql,params,1)
|
|
return data
|
|
|
|
def get_assets_users_db(self,uname):
|
|
if uname:
|
|
strsql = "select ID,uname,tellnum,tell_username from assets_user where uname like %s;"
|
|
params = (f'%{uname}%')
|
|
datas = self.safe_do_select(strsql,params)
|
|
else:
|
|
strsql = "select ID,uname,tellnum,tell_username from assets_user;"
|
|
datas = self.do_select(strsql)
|
|
return datas
|
|
|
|
def update_assets_users_db(self,IP,owner_id,itype):
|
|
if itype ==1:
|
|
strsql = '''
|
|
update ip_assets set owner_id = %s where ip_address = %s;
|
|
'''
|
|
else:
|
|
strsql = '''
|
|
update url_assets set owner_id = %s where ID = %s;
|
|
'''
|
|
params = (owner_id,IP)
|
|
bok,_ = self.safe_do_sql(strsql,params)
|
|
error = ""
|
|
if not bok:
|
|
error = "修改资产所属用户失败"
|
|
return bok,error
|
|
|
|
def get_port_latest_db(self,ip):
|
|
strsql = '''
|
|
select po.port,po.service,po.version,po.status from port_assets as po
|
|
left join ip_assets as ip
|
|
on ip.id = po.ip_id and ip.scan_count = po.scan_count
|
|
where ip.ip_address = %s
|
|
'''
|
|
params = (ip)
|
|
datas = self.safe_do_select(strsql,params)
|
|
return datas
|
|
|
|
def get_ip_url_latest_db(self,ip):
|
|
strsql = '''
|
|
select url.URL,url.subdomains,url.registrar,url.emails,url.creation_date,url.expiration_date from url_assets as url
|
|
left join ip_to_url as i2u on i2u.url_id = url.ID
|
|
left join ip_assets as ip on ip.id = i2u.ip_id
|
|
where ip_address = %s;
|
|
'''
|
|
params = (ip)
|
|
datas = self.safe_do_select(strsql,params)
|
|
return datas
|
|
|
|
def get_ip_url_history_db(self,ip):
|
|
#先把ip_id获取到
|
|
strsql = "select id from ip_assets where ip_address=%s;"
|
|
params = (ip,)
|
|
data = self.safe_do_select(strsql,params,1)
|
|
if not data:
|
|
return None
|
|
ip_id = data[0]
|
|
|
|
strsql = '''
|
|
select url.URL,i2u.create_time as time,'add' as type from ip_to_url as i2u
|
|
left join url_assets as url on url.ID = i2u.url_id
|
|
where ip_id=%s
|
|
union all
|
|
select url.URL,i2u_his.del_time as time,'del' as type from ip_to_url_his as i2u_his
|
|
left join url_assets as url on url.ID = i2u_his.url_id
|
|
where ip_id=%s
|
|
order by time DESC;
|
|
'''
|
|
params =(ip_id,ip_id)
|
|
datas = self.safe_do_select(strsql,params)
|
|
return datas
|
|
|
|
def del_url_assets_db(self,url_id):
|
|
#url assets
|
|
strsql ="delete from url_assets where ID=%s;"
|
|
params = (url_id,)
|
|
bok,_ = self.safe_do_sql(strsql,params)
|
|
#url_assets_his
|
|
strsql = "delete from url_assets_his where url_id=%s;"
|
|
bok, _ = self.safe_do_sql(strsql, params)
|
|
#ip_to_url
|
|
strsql = "delete from ip_to_url where url_id = %s;"
|
|
bok, _ = self.safe_do_sql(strsql, params)
|
|
#ip_to_url_his
|
|
strsql = "delete from ip_to_url_his where url_id = %s;"
|
|
bok, _ = self.safe_do_sql(strsql, params)
|
|
return True,"删除URL资产成功"
|
|
|
|
def get_last_task_by_ip(self,ip):
|
|
#寻找该IP最新完成的任务
|
|
strsql = '''
|
|
SELECT t.ID AS latest_task_id,t.start_time
|
|
FROM ip_assets AS ia
|
|
JOIN task_to_ip AS ti ON ia.id = ti.ip_id
|
|
JOIN task AS t ON ti.task_id = t.ID
|
|
WHERE ia.ip_address = %s
|
|
ORDER BY t.start_time DESC
|
|
LIMIT 1;
|
|
'''
|
|
params = (ip,)
|
|
data = self.safe_do_select(strsql ,params,1)
|
|
if data:
|
|
return data[0]
|
|
else:
|
|
return None
|
|
|
|
def del_ip_assets(self,ip):
|
|
#删除IP资产 -- 资产库中跟该资产相关的数据都需要删除,return bok,error
|
|
bok = False
|
|
# ip_assets
|
|
sql = """
|
|
DELETE FROM ip_assets
|
|
WHERE ip_address = %s
|
|
RETURNING id;
|
|
"""
|
|
params = (ip,)
|
|
bok,ip_id= self.safe_do_sql(sql,params,2)
|
|
if ip_id:
|
|
#task_to_ip
|
|
strsql = "delete from task_to_ip where ip_id = %s;"
|
|
bok,_ = self.safe_do_sql(strsql,params)
|
|
#ip_to_url
|
|
strsql = "delete from ip_to_url where ip_id = %s;"
|
|
bok, _ = self.safe_do_sql(strsql, params)
|
|
#ip_to_url_his
|
|
strsql = "delete from ip_to_url_his where ip_id = %s;"
|
|
bok, _ = self.safe_do_sql(strsql, params)
|
|
#port_assets
|
|
strsql = "delete from port_assets where ip_id = %s;"
|
|
bok, _ = self.safe_do_sql(strsql, params)
|
|
return True,"删除成功"
|
|
|
|
def get_url_assets_db(self,url,owner,email):
|
|
url_assets = []
|
|
strsql = '''
|
|
select url.ID,url.URL,au.uname,url.emails,url.update_time,url.expiration_date,ip.ip_count,url.Registrant,au.tellnum,au.tell_username,au.ID from url_assets as url
|
|
left join assets_user as au on au.ID = url.owner_id
|
|
left join (select url_id,count(*) as ip_count from ip_to_url group by url_id) as ip on url.ID = ip.url_id
|
|
'''
|
|
conditions = []
|
|
params = []
|
|
if url and url.strip():
|
|
conditions.append("url.URL like %s")
|
|
params.append(f"%{url}%")
|
|
if owner and owner.strip():
|
|
conditions.append("au.uname like %s")
|
|
params.append(f"%{owner}%")
|
|
if email and email.strip():
|
|
conditions.append("url.emails like %s")
|
|
params.append(f"%{email}%")
|
|
if len(conditions) > 0:
|
|
strsql += " WHERE " + " AND ".join(conditions)
|
|
# 执行查询（将参数转为元组）
|
|
url_assets = self.safe_do_select(strsql, tuple(params))
|
|
return url_assets
|
|
|
|
def get_url_to_ip_db(self,url_id):
|
|
strsql = '''
|
|
select ia.ip_address,itu.create_time from ip_assets ia
|
|
left join ip_to_url itu on itu.ip_id = ia.id
|
|
where itu.url_id = %s;
|
|
'''
|
|
params = (url_id,)
|
|
last_to_ips = self.safe_do_select(strsql,params)
|
|
strsql = '''
|
|
select ia.ip_address,itu.del_time from ip_assets ia
|
|
left join ip_to_url_his itu on itu.ip_id = ia.id
|
|
where itu.url_id = %s;
|
|
'''
|
|
params = (url_id,)
|
|
his_to_ips = self.safe_do_select(strsql, params)
|
|
return last_to_ips,his_to_ips
|
|
|
|
def get_owner_db(self,owner, owner_type, contact, tellnum):
|
|
strsql = '''
|
|
SELECT
|
|
u.ID,
|
|
u.itype,
|
|
u.uname,
|
|
u.tellnum,
|
|
u.tell_username,
|
|
u.ID_num,
|
|
IFNULL(ip.ip_count, 0) AS ip_count,
|
|
IFNULL(url.url_count, 0) AS url_count,
|
|
IFNULL(ip.ip_count, 0) + IFNULL(url.url_count, 0) AS total_assets
|
|
FROM assets_user u
|
|
LEFT JOIN (
|
|
SELECT owner_id, COUNT(*) AS ip_count
|
|
FROM ip_assets
|
|
GROUP BY owner_id
|
|
) ip ON u.ID = ip.owner_id
|
|
LEFT JOIN (
|
|
SELECT owner_id, COUNT(*) AS url_count
|
|
FROM url_assets
|
|
GROUP BY owner_id
|
|
) url ON u.ID = url.owner_id
|
|
|
|
'''
|
|
conditions = []
|
|
params = []
|
|
# 按需添加其他条件
|
|
if owner and owner.strip():
|
|
conditions.append("uname like %s")
|
|
params.append(f"%{owner}%")
|
|
|
|
if owner_type and owner_type.strip():
|
|
conditions.append("itype=%s")
|
|
params.append(owner_type)
|
|
|
|
if contact and contact.strip():
|
|
conditions.append("tell_username like %s")
|
|
params.append(f"%{contact}%")
|
|
|
|
if tellnum and tellnum.strip():
|
|
conditions.append("tellnum like %s")
|
|
params.append(f"%{tellnum}%")
|
|
|
|
# 组合完整的WHERE子句
|
|
if len(conditions) > 0:
|
|
strsql += " WHERE " + " AND ".join(conditions)
|
|
|
|
# 执行查询（将参数转为元组）
|
|
datas = self.safe_do_select(strsql, tuple(params))
|
|
return datas
|
|
|
|
def del_owner_db(self,id):
|
|
strsql = "delete from assets_user where ID=%s;"
|
|
params = (id,)
|
|
bok,_ = self.safe_do_sql(strsql,params)
|
|
#删除IP和URL与owner的关系
|
|
strsql = "update ip_assets set owner_id = 0 where owner_id = %s;"
|
|
bok,_ = self.safe_do_sql(strsql,params)
|
|
strsql = "update url_assets set owner_id = 0 where owner_id = %s;"
|
|
bok, _ = self.safe_do_sql(strsql, params)
|
|
return True,""
|
|
|
|
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()
|
|
|
|
#全局的单一实例
|
|
app_DBM = DBManager()
|
|
app_DBM.connect()
|
|
|
|
if __name__ == "__main__":
|
|
# mDBM = DBManager()
|
|
# mDBM.connect()
|
|
# print(mDBM.start_task("11","22"))
|
|
list_a = ['1','2','3','4','5','6']
|
|
str_a = ','.join(list_a)
|
|
print(str_a)
|