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 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 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 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 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_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 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 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": "创建任务时出现未知异常，请联系管理员查看异常日志"})