def check_url(task_id, task_access_key):
"""
接收来自浏览器的流量，对流量进行解析分类之后，存放到redis中，支持多个用户同时协作对一个任务进行测试
:param task_id:
:param task_access_key:
:return:
"""
from common.config_util import get_system_config
try:
post_data = request.get_json(force=True)
current_user_name = session["user_name"]
if TaskService.count(where=(Task.id == task_id,
Task.access_key == task_access_key)) == 0:
return jsonify(status=403,
message="发送失败",
data={
"extra_info":
"taskid或者accesskey不正确，插件请同步一次最新任务",
"site":
get_system_config()['front_end']['index']
})
tasks = TaskService.get_fields_by_where(
where=(Task.id == task_id, Task.access_key == task_access_key))
if tasks[0].task_status in [TaskStatus.DONE, TaskStatus.KILLED]:
return jsonify(status=400,
message="发送失败",
data={
"extra_info": "该任务已经结束，客户端请重新同步或者创建最新任务",
"site":
get_system_config()['front_end']['index']
})
if post_data is not None and "data" in post_data and RedisService.create_urlclassifications(
task_id, post_data):
raw_request_data = post_data.get('data', '{}')
scan_celery.delay(raw_request_data, task_id, current_user_name,
TaskStatus.NONE)
return jsonify(status=200,
message="发送成功",
data={"extra_info": "发送到后端扫描引擎成功"})
return jsonify(status=200,
message="发送失败",
data={"extra_info": "发送到后端引擎的数据不符合格式或者已经发送过"})
except Exception as e:
logger.exception("check_url exception")
return jsonify(status=500,
message="未知异常",
data={"extra_info": "创建任务时出现未知异常，请联系管理员查看异常日志"})
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
})