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))
