def stop_task():
"""
关闭任务，关闭任务之后将用户任务信息进行持久化到数据库，包括结束任务时间
:return:
"""
post_data = request.get_json(force=True)
if has_dict_value_blank(post_data, ["task_id"]):
return jsonify(status=400,
message="结束任务失败",
data={"extra_info": "task_id缺失,无法结束任务"})
try:
post_task_id = int(post_data.get("task_id"))
TaskService.update(fields=({
Task.task_status: TaskStatus.KILLED
}),
where=(Task.id == post_task_id))
RedisService.stop_task(post_task_id)
RedisService.clean_urlclassifications(post_task_id)
return jsonify(status=200,
message="结束任务成功",
data={"extra_info": "该任务由管理员关闭"})
except Exception as e:
logger.exception("stop_task exception")
return jsonify(status=500,
message="未知异常",
data={"extra_info": "创建任务时出现未知异常，请联系管理员查看异常日志"})
def create_task():
"""
创建任务，可以由浏览器插件或者直接在平台上创建，redis缓存查询用户信息
:return:
"""
try:
post_data = request.get_json(force=True)
if json_utils.has_dict_value_blank(
post_data, ['hook_rule', 'read_agreement', 'task_name']):
return jsonify(status=400,
message="创建任务失败",
data={"extra_info": "新建任务时没有设置网址正则或任务名称"})
if not post_data.get("read_agreement"):
return jsonify(status=400,
message="创建任务失败",
data={"extra_info": "请阅读用户协议并点击同意"})
current_user_name = session["user_name"]
post_hook_rule = post_data.get("hook_rule")
post_task_name = post_data.get("task_name")
current_user = RedisService.get_user(user_name=current_user_name)
receivers_email = current_user[
"email"] if "email" in current_user else None
task = TaskService.save(create_user_id=current_user["id"],
task_name=post_task_name,
receivers_email=receivers_email,
hook_rule=post_hook_rule)
#if UserTaskService.count(where=(UserTask.task_id == task.id, UserTask.user_id == current_user["id"])) == 0:
UserTaskService.save(task_id=task.id, user_id=current_user["id"])
RedisService.create_task(task.id, post_hook_rule, current_user_name,
TaskStatus.WORKING)
RedisService.update_user_field(current_user_name, "current_task_id",
task.id)
UserService.update(fields=({
User.recent_operation_time:
datetime.datetime.now()
}),
where=(User.user_name == current_user_name))
task_access_key = generate_access_key(
task.id, current_user_name).decode("utf-8")
TaskService.update(fields=({
Task.access_key: task_access_key
}),
where=(Task.id == task.id))
scan_celery.delay(post_data, task.id, current_user_name,
TaskStatus.WORKING)
return jsonify(status=200,
message="创建任务成功",
data={
"task_id": task.id,
"full_name": current_user["full_name"],
"create_time": get_current_time(),
"task_access_key": task_access_key
})
except Exception as e:
logger.exception("create_task exception")
return jsonify(status=500,
message="未知异常",
data={"extra_info": "创建任务时出现未知异常，请联系管理员查看异常日志"})
def update_task():
"""
更新任务信息，最主要的是修改hook_url
:return:
"""
try:
put_data = request.get_json(force=True)
if has_dict_value_blank(put_data, ["hook_rule", "task_id"]):
return jsonify(status=400,
message="更新任务失败",
data={"extra_info": "请确认是否正确传入hook_rule,task_id参数"})
hook_rule = put_data.get("hook_rule")
task_id = put_data.get("task_id")
# Task表修改一下hook_rule
TaskService.update(fields=({
Task.hook_rule: hook_rule
}),
where=(Task.id == task_id))
# redis更改任务hook_rule
RedisService.update_task_hook_rule(task_id, hook_rule)
return jsonify(status=200,
message="更新任务成功",
data={"extra_info": "修改成功,该任务由管理员修改"})
except Exception as e:
logger.exception("update_task exception")
return jsonify(status=500,
message="未知异常",
data={"extra_info": "创建任务时出现未知异常，请联系管理员查看异常日志"})
def scan_celery(package, task_id, create_user, status):
"""
celey 调度模式扫描
注意，scan_celery函数只有在celey开启的时候才会
:param package:
:param task_id:
:param create_user:
:param status:
:return:
"""
logger.setLevel(logging.INFO)
logger.info(logo)
if status == TaskStatus.WORKING:
# 更新任务状态和hunter状态
current_task_status = TaskService.get_task_status(task_id=task_id)
if current_task_status and current_task_status < TaskStatus.WORKING:
TaskService.update(fields=({
Task.task_status: TaskStatus.WORKING
}),
where=(Task.id == task_id))
TaskService.update(fields=({
Task.xssfork_status: TaskStatus.WORKING
}),
where=(Task.id == task_id))
logger.warn(
"there is a task [task_id:{}, create_user:{}] has start".format(
task_id, create_user))
elif status == TaskStatus.KILLED:
try:
TaskService.update(fields=({
Task.xssfork_status: TaskStatus.DONE
}),
where=(Task.id == task_id))
current_task = TaskService.get_fields_by_where(
where=(Task.id == task_id))[0]
if current_task.hunter_status == TaskStatus.DONE and current_task.sqlmap_status == TaskStatus.DONE \
and current_task.xssfork_status == TaskStatus.DONE:
TaskService.update(fields=({
Task.task_status: TaskStatus.DONE
}),
where=(Task.id == task_id))
task_notice_celery.delay(
message={
"type": BroadCastType.TASK,
"action": BroadCastAction.COMPLETE_TASK_NOTIFICATION,
"data": {
"task_id": task_id
}
})
except Exception:
logger.exception("scan_celery error")
logger.warn(
"there is a task [task_id:{}, create_user:{}] has killed".format(
task_id, create_user))
else:
scan(package=package,
task_id=task_id,
create_user=create_user,
status=status)
def testGetTasksUrlNum(self):
"""
测试执行SQL
:return:
"""
from model.task import TaskService
import logging
logger = logging.getLogger('peewee')
logger.addHandler(logging.StreamHandler())
logger.setLevel(logging.DEBUG)
TaskService.get_tasks_url_num(task_id=1, task_status=3)
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 show_task_spent_time():
"""
展示每个任务从创建到结束的时间
:return:
"""
try:
current_user_name = session.get('user_name')
current_user = RedisService.get_user(current_user_name)
count = int(request.args.get("count"))
tasks = TaskService.get_tasks(user_id=current_user.id)
tasks = tasks[-count:]
tasks_json = [{
"id":
task.id,
"task_name":
task.task_name,
"created_time":
task.created_time.strftime("%Y-%m-%d %H:%M")
if task.created_time else "",
"killed_time":
task.killed_time.strftime("%Y-%m-%d %H:%M")
if task.killed_time else "",
"spend_time":
minutes_(task.killed_time, task.created_time)
} for task in tasks]
return jsonify(status=200, message="查询成功", data=tasks_json)
except Exception as e:
logger.exception("show_task_spent_time rasie error")
if isinstance(e, KeyError):
return jsonify(status=400, message="查询失败", reason="请传递参数count")
return jsonify(status=500,
message="未知异常",
data={"extra_info": "创建任务时出现未知异常，请联系管理员查看异常日志"})
def show_scan_records():
try:
current_user_name = session.get('user_name')
current_user = RedisService.get_user(current_user_name)
tasks = TaskService.get_tasks_url_vuln_num(user_id=current_user.id)
response_data = list()
for task in tasks:
risk_level = VulnerabilityService.get_risk_level(task)
response_data.append({
"task_name":
task.task_name,
"created_time":
task.created_time.strftime("%Y-%m-%d %H:%M"),
"task_id":
task.id,
"urls_num":
task.urls_num,
"vulns_num":
task.vulns_num,
"risk_level":
risk_level
})
response_data.reverse()
return jsonify(status=200, message="查询成功", data=response_data)
except Exception:
logger.exception("show_current_tasks rasie error")
return jsonify(status=500,
message="未知异常",
data={"extra_info": "创建任务时出现未知异常，请联系管理员查看异常日志"})
def testGetCurrentTasks(self):
"""
获取最新任务
:return:
"""
import logging
logger = logging.getLogger('peewee')
logger.addHandler(logging.StreamHandler())
logger.setLevel(logging.DEBUG)
from model.task import TaskService
task = TaskService.get_current_tasks(1)[0]
print(task)
def update_task():
"""
更新任务信息，最主要的是修改hook_url
:return:
"""
post_data = request.get_json(force=True)
if json_utils.has_dict_value_blank(post_data, ["hook_rule", "task_id"]):
return jsonify(status=400,
message="更新任务失败",
data={"extra_info": "请确认是否正确传入hook_rule,task_id参数"})
try:
post_hook_rule = post_data.get("hook_rule")
post_task_id = post_data.get("task_id")
current_user_name = session["user_name"]
# current_user_name = "b5mali4"
current_user = RedisService.get_user(current_user_name)
if UserTaskService.count(
where=(UserTask.user_id == int(current_user["id"]),
UserTask.task_id == post_task_id)) == 0:
return jsonify(status=403,
message="更新任务失败",
data={"extra_info": "请勿尝试非法更改非自己权限的任务"})
TaskService.update(fields=({
Task.hook_rule: post_hook_rule
}),
where=(Task.id == post_task_id))
RedisService.update_task_hook_rule(post_task_id, post_hook_rule)
return jsonify(status=200,
message="更新任务成功",
data={"extra_info": "修改成功,刷新页面即可看到更改结果"})
except Exception as e:
logger.exception("update_task exception")
if isinstance(e, UserNotFoundInRedisException):
return jsonify(status=403,
message="结束任务失败",
data={"extra_info": "请勿尝试非法关闭非自己权限任务"})
return jsonify(status=500,
message="未知异常",
data={"extra_info": "创建任务时出现未知异常，请联系管理员查看异常日志"})
def testGetFrontData(self):
"""
获取前几天
:return:
"""
import logging
logger = logging.getLogger('peewee')
logger.addHandler(logging.StreamHandler())
logger.setLevel(logging.DEBUG)
from common.system_util import get_front_date
from model.task import Task, TaskService
front_date = get_front_date(100)
tasks = TaskService.get_fields_by_where(
where=(Task.created_time >= front_date))
for task in tasks:
print(task)
def generate_report(self, task_id):
"""
生成邮件发送报告
:param cls:
:param task_id:
:return:
"""
current_task = TaskService.get_fields_by_where(where=(Task.id == task_id))[0]
vulns_info = VulnerabilityService.get_fields_by_where(where=(Vulnerability.task_id == task_id))
users = UserService.get_users(task_id=task_id)
if len(vulns_info) <= 0:
content = """<br>你好，欢迎使用Hunter，本次扫描结束，扫描到你有0个漏洞。详情请可登录{}查看<br>""".format(
get_system_config()['front_end']['index'])
else:
content = """<br>你好，欢迎使用Hunter，本次扫描结束，扫描到你有{}个漏洞。任务预览如下,详情请登录{}查看<br>""".format(len(vulns_info),
get_system_config()[
'front_end'][
'index'])
content += """
<table frame='border' cellpadding='15' cellspacing='0' align='center' style='border: 1px solid #d6d3d3;'>
<tr style='background: #e6e6e6;'>
<th style="border-right: 1px solid #bfbfbf;">序号</th>
<th style="border-right: 1px solid #bfbfbf;">漏洞等级</th>
<th style="border-right: 1px solid #bfbfbf;">漏洞类型</th>
<th style="border-right: 1px solid #bfbfbf;">漏洞详情</th>
</tr>
"""
index = 0
for vuln_info in vulns_info:
index += 1
vuln_detail_url = '<a href="{}">{}</a>'.format(
get_system_config()['front_end']['vuln_route'] + str(task_id),
vuln_info.info)
content += """
<tr>
<td style="border-right: 1px solid #bfbfbf;">{}</td>
<td style="border-right: 1px solid #bfbfbf;">{}</td>
<td style="border-right: 1px solid #bfbfbf;">{}</td>
<td style="border-right: 1px solid #bfbfbf;">{}</td>
</tr>
""".format(index, vuln_info.level, vuln_info.chinese_type, vuln_detail_url)
content += """</table>"""
return content, ",".join([user.email for user in users if user.email])
def get_working_task_info():
"""
根据用户名获取得到任务信息，用于openrestry代理模块取得最新未结束的任务，接口无需鉴权
:param user_name:
:return:
"""
try:
user_name = request.args.get("user_name")
user = RedisService.get_user(user_name)
tasks = TaskService.get_working_tasks(user.id)
descriptions = ["该任务处于激活状态", "该任务正在进行扫描", "该任务已经扫描完成", "该任务已被用户关闭"]
if tasks is None or len(tasks) == 0:
return jsonify(status=200,
message="查询成功",
data={
"id": -1,
"status": TaskStatus.NONE,
"description": "不存在正在工作的任务，请登录平台先创建任务"
})
current_working_task = tasks[0]
return jsonify(status=200,
message="查询成功",
data={
"id":
current_working_task.id,
"status":
current_working_task.task_status,
"description":
descriptions[current_working_task.task_status]
})
except Exception as e:
logger.exception("get_task_info raise error")
if isinstance(e, KeyError):
return jsonify(status=400,
message="查询失败",
data={"extra_info": "未传递参数user_name"})
elif isinstance(e, UserNotFoundInRedisException):
return jsonify(status=400,
message="查询失败",
data={"extra_info": "未找到该用户信息"})
return jsonify(status=500,
message="未知异常",
data={"extra_info": "发生未知异常，请联系管理员查看异常日志"})
def list_tasks_by_num():
"""
获取系统中的任务列表，用于数据大盘展示用户使用最新动态，需要跨表连接获取数据
SELECT * FROM user INNER JOIN usertask ON usertask.user_id = user.id where usertask.task_id = 2333
:return:
"""
try:
num = int(request.args.get("num"))
tasks_cursor_wrapper = TaskService.get_fields_by_where(
fields=(Task.created_time, Task.killed_time, Task.task_name,
Task.id))
tasks_list = OrmModelJsonSerializer.serializer(tasks_cursor_wrapper)
tasks_list.reverse()
if len(tasks_list) > num:
tasks_result = tasks_list[0:num]
else:
tasks_result = tasks_list
result = list()
for task_result in tasks_result:
user = UserService.get_users(task_id=task_result["id"])[0]
task_result["dept_name"] = user.dept_name
task_result["full_name"] = user.full_name
task_result["user_name"] = user.user_name
result.append(task_result)
response_data = jsonify(status=200,
message="查询成功",
data={
'tasks': result,
'num': len(result)
})
except Exception as e:
logger.exception("list_tasks raise error")
return jsonify(status=200,
message="查询成功",
data={
'tasks': [],
'num': 0
})
return response_data
def list_scan_record():
"""
根据task_id查询扫描记录
:return:
"""
try:
user_id = request.args.get("user_id")
tasks = TaskService.get_tasks_url_vuln_num(user_id=user_id)
response_data = list()
for task in tasks:
risk_level = VulnerabilityService.get_risk_level(task)
response_data.append({
"task_name":
task.task_name,
"created_time":
task.created_time.strftime("%Y-%m-%d %H:%M"),
"task_id":
task.id,
"urls_num":
task.urls_num,
"vulns_num":
task.vulns_num,
"risk_level":
risk_level
})
response_data.reverse()
return jsonify(status=200, message="查询成功", data=response_data)
except Exception as e:
logger.exception("list_scan_record raise error")
if isinstance(e, KeyError):
return jsonify(status=400,
message="查询失败",
data={"extra_info": "未传递taskid"})
return jsonify(status=500,
message="未知异常",
data={"extra_info": "发生未知异常，请联系管理员查看异常日志"})
def list_task_spent_time():
"""
列出最近n次任务扫描消耗的时间，用于数据大盘第四个图
:return:
"""
from common.system_util import minutes_
try:
count = int(request.args.get("count"))
response_data = list()
tasks = TaskService.get_fields_by_where(
fields=(Task.id, Task.task_name, Task.created_time,
Task.killed_time))[-count:]
for task in tasks:
created_time = task.created_time.strftime(
"%Y-%m-%d %H:%M") if task.created_time else ""
killed_time = task.killed_time.strftime(
"%Y-%m-%d %H:%M") if task.killed_time else ""
response_data.append({
"id":
task.id,
"task_name":
task.task_name,
"created_time":
created_time,
"killed_time":
killed_time,
"spend_time":
minutes_(task.killed_time, task.created_time)
})
return jsonify(status=200, message="查询成功", data=response_data)
except Exception as e:
logger.exception("list_task_spent_time raise error")
if isinstance(e, KeyError):
return jsonify(status=400,
message="查询失败",
data={"extra_info": "请传递参数count"})
return jsonify(status=500,
message="未知异常",
data={"extra_info": "发生未知异常，请联系管理员查看异常日志"})
def testGetTasksUrlVulnNum(self):
from model.task import TaskService
from model.hunter_model import OrmModelJsonSerializer
import logging
logger = logging.getLogger('peewee')
logger.addHandler(logging.StreamHandler())
logger.setLevel(logging.DEBUG)
tasks = TaskService.get_tasks_url_vuln_num(user_id=1)
response_data = [{
"task_name":
task.task_name,
"created_time":
task.created_time.strftime("%Y-%m-%d %H:%M"),
"task_id":
task.id,
"urls_num":
task.urls_num,
"vulns_num":
task.vulns_num
} for task in tasks]
for response in response_data:
print(response)
def modify_user_info_cache_session(user_name, db_user):
"""
认证成功之后，修改redis中用户信息并设置session
:return:
"""
# 存入到redis，和数据库中的数据一致，用户id是最关键的
RedisService.update_user(
user_name, {
"id": db_user.id,
"user_name": db_user.user_name,
"full_name": db_user.full_name,
"dept_name": db_user.dept_name,
"role": db_user.role,
"mobile_phone": db_user.mobile_phone,
"email": db_user.email
})
try:
current_task = TaskService.get_working_tasks(user_id=db_user.id)[0]
RedisService.update_user(user_name,
{"current_task_id": current_task.id})
except IndexError:
pass
# 设置session
session["user_name"] = user_name
def show_current_task():
"""
显示当前任务正在运行的任务
:return:
"""
try:
current_user_name = session["user_name"]
current_user = RedisService.get_user(current_user_name)
current_task = TaskService.get_working_tasks(
user_id=current_user.id)[0]
if current_task:
hook_rule = RedisService.get_task(current_task.id)["hook_rule"]
unscaned_url_num = UrlService.count(
where=(Url.task_id == current_task.id,
Url.status != TaskStatus.DONE))
scaned_url_num = UrlService.count(
where=(Url.task_id == current_task.id,
Url.status == TaskStatus.DONE))
total_url_num = unscaned_url_num + scaned_url_num
if current_task.task_status in [
TaskStatus.KILLED, TaskStatus.DONE
]:
percent = 100
else:
percent = 0 if total_url_num == 0 else (scaned_url_num /
total_url_num) * 100
response_data = jsonify(
status=200,
message="查询成功",
data={
'receiver_emails':
current_task.receivers_email,
'task_name':
current_task.task_name,
'create_time':
current_task.created_time.strftime("%Y-%m-%d %H:%M"),
'percent':
percent,
'unscaned_url_num':
unscaned_url_num,
'scaned_url_num':
scaned_url_num,
'total_url_num':
total_url_num,
'hook_rule':
hook_rule,
'task_id':
current_task.id,
"task_access_key":
current_task.access_key,
'task_status':
current_task.task_status,
'user_name':
current_user_name
})
return response_data
except Exception as e:
if isinstance(e, IndexError):
return jsonify(status=400,
message="获取失败",
data={"extra_info": "后台无正在运行任务，请登录后台并创建任务"})
logger.exception("show_current_task rasie error")
return jsonify(status=500,
message="获取失败",
data={"extra_info": "未知异常，可以联系管理员到后台查看"})
def stop_task():
"""
关闭任务，关闭任务之后将用户任务信息进行持久化到数据库，包括结束任务时间
:return:
"""
post_data = request.get_json(force=True)
if json_utils.has_dict_value_blank(post_data, ["task_id"]):
return jsonify(status=400,
message="结束任务失败",
data={"extra_info": "task_id缺失,无法结束任务"})
post_task_id = int(post_data.get("task_id"))
current_user_name = session["user_name"]
try:
user = RedisService.get_user(current_user_name)
if UserTaskService.count(
where=(UserTask.user_id == user["id"],
UserTask.task_id == post_task_id)) == 0:
return jsonify(status=403,
message="结束任务失败",
data={"extra_info": "请勿尝试非法关闭非自己权限任务"})
task = TaskService.get_fields_by_where(
fields=(Task.task_status), where=(Task.id == post_task_id))[0]
if task.task_status in [TaskStatus.DONE, TaskStatus.KILLED]:
return jsonify(status=200,
message="结束任务成功",
data={
"fullname": user["full_name"],
"extra_info": "该任务早已经结束,请登录后台查看扫描结果",
"stop_time": get_current_time()
})
TaskService.update(fields=({
Task.task_status: TaskStatus.KILLED,
Task.killed_time: datetime.datetime.now()
}),
where=(Task.id == post_task_id))
UserService.update(fields=({
User.recent_operation_time:
datetime.datetime.now()
}),
where=(User.user_name == current_user_name))
scan_celery.delay(post_data, post_task_id, current_user_name,
TaskStatus.KILLED)
RedisService.stop_task(post_task_id)
RedisService.clean_urlclassifications(post_task_id)
return jsonify(status=200,
message="结束任务成功",
data={
"full_name":
user["full_name"],
"extra_info":
"请登录后台查看扫描结果",
"stop_time":
datetime.datetime.now().strftime("%Y-%m-%d %H:%M")
})
except Exception as e:
logger.exception("stop_task exception")
if isinstance(e, UserNotFoundInRedisException):
return jsonify(status=403,
message="结束任务失败",
data={
"extra_info": "认证失败,请重新登录进行授权",
"auth_site": ""
})
return jsonify(status=500,
message="未知异常",
data={"extra_info": "创建任务时出现未知异常，请联系管理员查看异常日志"})
def show_current_tasks():
"""
显示当前所有的任务列表，在响应中返回的结果为task_list和已经扫描的和未扫描的任务数目
:return:
"""
try:
working_tasks = list()
completed_tasks = list()
working_task_info_list = list()
current_user_name = session["user_name"]
current_user = RedisService.get_user(current_user_name)
tasks = TaskService.get_tasks_url_vuln_num(user_id=current_user.id)
for task in tasks:
if task.task_status <= TaskStatus.WORKING:
working_tasks.append(task)
if task.task_status == TaskStatus.DONE:
completed_tasks.append(task)
for working_task in working_tasks:
hook_rule = RedisService.get_task(working_task.id)["hook_rule"]
unscaned_url_num = working_task.unscaned_urls_num
scaned_url_num = working_task.scaned_urls_num
total_url_num = unscaned_url_num + scaned_url_num
if working_task.task_status in [
TaskStatus.KILLED, TaskStatus.DONE
]:
percent = 100
else:
percent = 0 if total_url_num == 0 else int(
(scaned_url_num / total_url_num) * 100)
working_task_info_list.append({
'receiver_emails':
working_task.receivers_email,
'task_name':
working_task.task_name,
'create_time':
working_task.created_time.strftime("%Y-%m-%d %H:%M"),
'percent':
percent,
'unscaned_url_num':
unscaned_url_num,
'scaned_url_num':
scaned_url_num,
'total_url_num':
total_url_num,
'hook_rule':
hook_rule,
'task_id':
working_task.id,
"task_access_key":
working_task.access_key,
'task_status':
working_task.task_status
})
response = jsonify(status=200,
message="查询成功",
data={
"working_task_info_list":
working_task_info_list,
"working_task_num":
len(working_tasks) -
1 if len(working_tasks) > 0 else 0,
"completed_task_num":
len(completed_tasks)
})
return response
except Exception:
logger.exception("show_current_tasks rasie error")
return jsonify(status=500,
message="未知异常",
data={"extra_info": "创建任务时出现未知异常，请联系管理员查看异常日志"})
def list_tasks():
"""
显示所有任务列表，方便管理任务
:return:
"""
try:
task_id = request.args.get("task_id")
task_status = request.args.get('status')
# 构造条件查询元组
task_info_list = list()
tasks = TaskService.get_tasks_url_num(task_id=task_id,
task_status=task_status)
for task in tasks:
hook_rule = task.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 int(
(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,
"create_user_name":
task.create_user_name
})
task_info_list.reverse()
response = jsonify(status=200, message="查询成功", data=task_info_list)
return response
except Exception as e:
logger.exception("show_current_tasks rasie error")
if isinstance(e, BaseHunterException):
return jsonify(status=400,
message=str(e),
data={"extra_info": "查询任务时传入非法的task_id"})
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
})