def get_create_user(task_id): """ 根据任务ID获取到创建任务者 SELECT * FROM user WHERE id = (SELECT create_user_id FROM task WHERE task.id = task_id) :return: """ users = User.select().where(User.id == (Task.select(Task.create_user_id).where(Task.id == task_id))).execute() if len(users) == 1: return users[0] return None
def get_tasks_urls_vulns_num_by_days(user_id=None, day_range=None): """ A情况:day_range传入,user_id不传入 sql语句:SELECT *,(SELECT COUNT(*) FROM url WHERE url.task_id=task.id) AS url_num,(SELECT COUNT(*) FROM vulnerability WHERE vulnerability.task_id=task.id) AS vuln_num FROM task WHERE task.created_time > '1745-11-15' B情况:day_range传入,user_id传入 sql语句:SELECT *,(SELECT COUNT(*) FROM url WHERE url.task_id=task.id) AS url_num,(SELECT COUNT(*) FROM vulnerability WHERE vulnerability.task_id=task.id) AS vuln_num FROM task INNER JOIN usertask ON task.id = usertask.task_id WHERE task.created_time >= '1745-11-15' AND usertask.user_id = 1 C情况:day_range不传入,user_id不传入 sql语句:SELECT *,(SELECT COUNT(*) FROM url WHERE url.task_id=task.id) AS url_num,(SELECT COUNT(*) FROM vulnerability WHERE vulnerability.task_id=task.id) AS vuln_num FROM task D情况:day_range不传入,user_id传入 sql语句:SELECT *,(SELECT COUNT(*) FROM url WHERE url.task_id=task.id) AS url_num,(SELECT COUNT(*) FROM vulnerability WHERE vulnerability.task_id=task.id) AS vuln_num FROM task INNER JOIN usertask ON task.id = usertask.task_id WHERE usertask.user_id = 1 统计最近几天内的任务,url,漏洞数量 SELECT * FROM task WHERE DATE_SUB(CURDATE(), INTERVAL 10 DAY) <= created_time SELECT * FROM task INNER JOIN usertask ON usertask.task_id = task.id where usertask.user_id='1222' :param user_id: :param day: :return: """ from model.task import TaskService, Task from model.user_task import UserTask, UserTaskService from model.url import Url, UrlService from common.system_util import get_front_date from common.json_utils import dict_auto_add created_date2task_num = dict() created_date2urls_num = dict() created_date2vulns_num = dict() if user_id: task_total_num = TaskService.count(where=(Task.create_user_id == user_id)) url_total_num = Url.select(fn.COUNT(Url.id).alias('urls_total_num')).join(UserTask, JOIN.INNER, on=( UserTask.task_id == Url.task_id)).where(UserTask.user_id == user_id).execute()[0].urls_total_num vuln_total_num = \ Vulnerability.select(fn.COUNT(Vulnerability.id).alias('vulns_total_num')).join(UserTask, JOIN.INNER, on=( UserTask.task_id == Vulnerability.task_id)).where( UserTask.user_id == user_id).execute()[0].vulns_total_num else: task_total_num = TaskService.count() url_total_num = UrlService.count() vuln_total_num = VulnerabilityService.count() if day_range: front_date = get_front_date(day_range=day_range) if not user_id: # 情况A matched_tasks = Task.select(Task.id, Task.created_time, Url.select(fn.COUNT(Url.id)).alias('urls_num').where( Url.task_id == Task.id), Vulnerability.select(fn.COUNT(Vulnerability.id)).alias( 'vulns_num').where( Vulnerability.task_id == Task.id)).where( Task.created_time >= front_date).execute() else: # 情况B matched_tasks = Task.select(Task.id, Task.created_time, Url.select(fn.COUNT(Url.id)).alias('urls_num').where( Url.task_id == Task.id), Vulnerability.select(fn.COUNT(Vulnerability.id)).alias( 'vulns_num').where( Vulnerability.task_id == Task.id)).join(UserTask, JOIN.INNER, on=( Task.id == UserTask.task_id)).where(Task.created_time >= front_date, UserTask.user_id == user_id).execute() else: if not user_id: # 情况C matched_tasks = Task.select(Task.id, Task.created_time, Url.select(fn.COUNT(Url.id)).alias('urls_num').where( Url.task_id == Task.id), Vulnerability.select(fn.COUNT(Vulnerability.id)).alias( 'vulns_num').where( Vulnerability.task_id == Task.id)).execute() else: # 情况D matched_tasks = Task.select(Task.id, Task.created_time, Url.select(fn.COUNT(Url.id)).alias('urls_num').where( Url.task_id == Task.id), Vulnerability.select(fn.COUNT(Vulnerability.id)).alias( 'vulns_num').where( Vulnerability.task_id == Task.id)).join(UserTask, JOIN.INNER, on=( Task.id == UserTask.task_id, UserTask.user_id == user_id)).execute() for matched_task in matched_tasks: create_day = str(matched_task.created_time).split(" ")[0] dict_auto_add(created_date2task_num, create_day) dict_auto_add(created_date2urls_num, create_day, matched_task.urls_num) dict_auto_add(created_date2vulns_num, create_day, matched_task.vulns_num) return {"task": created_date2task_num, "task_total_num": task_total_num, "url": created_date2urls_num, "url_total_num": url_total_num, "vuln": created_date2vulns_num, "vuln_total_num": vuln_total_num}
def testGetTasksUrlsVulnsNumByDays(self): """ 测试get_tasks_urls_vulns_num_by_days函数耗时操作,优化sql语句,查看具体的索引使用情况 统计最近几天内的任务,url,漏洞数量 SELECT * FROM task WHERE DATE_SUB(CURDATE(), INTERVAL 10 DAY) <= created_time SELECT * FROM task INNER JOIN usertask ON usertask.task_id = task.id where usertask.user_id='1222' :return: """ import time from model.task import TaskService, Task from model.user_task import UserTask, UserTaskService from model.url import Url, UrlService from model.vulnerability import VulnerabilityService, Vulnerability from common.system_util import get_front_date from common.json_utils import dict_auto_add user_id = 1 day_range = 100000 tasks_num = {} urls_num = {} vulns_num = {} start_time = time.time() if user_id: task_total_num = TaskService.count( where=(Task.create_user_id == user_id)) else: task_total_num = TaskService.count() url_total_num = UrlService.count() vuln_total_num = VulnerabilityService.count() print("总共消耗了{}".format(time.time() - start_time)) if day_range: front_date = get_front_date(day_range=day_range) if user_id: matched_tasks = Task.select().join( UserTask, JOIN.INNER, on=(Task.id == UserTask.task_id)).where( UserTask.user_id == user_id, Task.created_time >= front_date).execute() else: matched_tasks = Task.select().where( Task.created_time >= front_date).execute() else: matched_tasks = TaskService.get_fields_by_where() for matched_task in matched_tasks: create_day = str(matched_task.created_time).split(" ")[0] dict_auto_add(tasks_num, create_day) import logging logger = logging.getLogger('peewee') logger.addHandler(logging.StreamHandler()) logger.setLevel(logging.DEBUG) tmp_url_num = UrlService.count( where=(Url.task_id == matched_task.id)) tmp_vuln_num = VulnerabilityService.count( where=(Vulnerability.task_id == matched_task.id)) dict_auto_add(urls_num, create_day, tmp_url_num) dict_auto_add(vulns_num, create_day, tmp_vuln_num) print("总共消耗了{}".format(time.time() - start_time)) print({ 'task': len(matched_tasks), 'task_total_num': task_total_num, 'url': urls_num, 'url_total_num': url_total_num, 'vuln': vulns_num, 'vuln_total_num': vuln_total_num })
def testListTasksSpendTime(self): """ 测试 http://10.211.55.2:8888/api/v1/admin/task/ api耗时 :return: """ import time import peewee from model.task import TaskService, Task from model.url import UrlService, Url from model.default_value import TaskStatus from api.service.redis_service import RedisService """ 测试 ListTask 耗时,便于对数据索引作出优化 :return: """ task_id = None status = None # 构造条件查询元组 query = list() if task_id is not None and task_id != "": query.append(Task.id == int(task_id)) if status is not None and status != "": query.append(Task.task_status == int(status)) # EXPLAIN SELECT *,(SELECT COUNT(*) FROM url WHERE url.task_id = task.id AND url.`status` !=2) AS 'unscaned_url_num',(SELECT COUNT(*) FROM url WHERE url.task_id = task.id AND url.`status` =2) AS 'scaned_urls_num' FROM task if len(query) > 0: tasks = Task.select(Task.receivers_email, Task.task_name, Task.created_time, Task.id, Task.access_key, Task.task_status, Url.select(fn.COUNT(Url.id)).alias('unscaned_urls_num').where(Url.task_id == Task.id, Url.status != TaskStatus.DONE), Url.select(fn.COUNT(Url.id)).alias('scaned_urls_num').where(Url.task_id == Task.id, Url.status == TaskStatus.DONE)).where( *tuple(query)).execute() else: import logging logger = logging.getLogger('peewee') logger.addHandler(logging.StreamHandler()) logger.setLevel(logging.DEBUG) tasks = Task.select(Task.receivers_email, Task.task_name, Task.created_time, Task.id, Task.access_key, Task.task_status, Url.select(fn.COUNT(Url.id)).alias('unscaned_urls_num').where(Url.task_id == Task.id, Url.status != TaskStatus.DONE), Url.select(fn.COUNT(Url.id)).alias('scaned_urls_num').where(Url.task_id == Task.id, Url.status == TaskStatus.DONE)).execute() task_info_list = list() for task in tasks: hook_rule = RedisService.get_task(task.id)["hook_rule"] unscaned_urls_num = task.unscaned_urls_num scaned_urls_num = task.scaned_urls_num total_url_num = unscaned_urls_num + scaned_urls_num if task.task_status in [TaskStatus.KILLED, TaskStatus.DONE]: percent = 100 else: percent = 0 if total_url_num == 0 else (scaned_urls_num / total_url_num) * 100 task_info_list.append({'receiver_emails': task.receivers_email, 'task_name': task.task_name, 'create_time': task.created_time.strftime("%Y-%m-%d %H:%M"), 'percent': percent, 'unscaned_url_num': unscaned_urls_num, 'scaned_url_num': scaned_urls_num, 'total_url_num': total_url_num, 'hook_rule': hook_rule, 'task_id': task.id, 'task_access_key': task.access_key, 'task_status': task.task_status}) task_info_list.reverse() print(task_info_list)
def get(self): type_ = self.get_argument('type') li = Task.select().where(Task.type == type_).order_by(Task.index) li = [o.to_dict() for o in li] self.finish(dict(data=li))