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 modify_plugin():
"""
禁用指定的插件，
服务端:1.从CHECKER_INSTANCE_DICT指定插件设置disable 2.并向MQ中发送一条消息
引擎端:1.等到引擎消费到指令从CHECKER_INSTANCE_DICT，也将该插件的disable设置为True
:return:
"""
from plugins.base.vuln_enum import PluginSwith
from common.plugins_util import modify_default_checkers
post_data = request.get_json(force=True)
param_list = ["name", "switch"]
if has_dict_value_blank(post_data, param_list):
return jsonify(
status=400,
message="更新失败",
data={"extra_info": "请保证%s任一参数值不为空" % ','.join(param_list)})
try:
post_checker_name = post_data.get("name")
post_switch = post_data.get("switch")
plugin_switch = PluginSwith.ON if post_switch else PluginSwith.OFF
checkers = load_default_checkers()
if post_checker_name not in checkers:
return jsonify(
status=400,
message="更新失败",
data={"extra_info": "不存在名为%s的插件" % post_checker_name})
PluginInfoService.update(
fields=({
PluginInfo.useable: plugin_switch
}),
where=(PluginInfo.plugin_name == post_checker_name))
RedisService.modify_plugin_switch(checkers[post_checker_name],
plugin_switch)
system_notice_celery.delay(
broadcast={
"type": BroadCastType.PLUGIN,
"action": BroadCastAction.MODIFY_PLUGIN,
"data": {
"name": post_checker_name,
"switch": plugin_switch
}
})
modify_default_checkers()
return jsonify(status=200,
message="更新成功",
data={
"extra_info":
"{}插件{}".format("启用" if plugin_switch else "禁用",
post_checker_name)
})
except Exception as e:
logger.exception("modify_plugin raise error")
return jsonify(status=500,
message="未知异常",
data={"extra_info": "修改插件信息时出现未知异常，请联系管理员查看异常日志"})
def testGetPluginConfigFromRedis(self):
"""
测试从redis中加载配置，方便测试初始化redis
:return:
"""
from api.service.redis_service import RedisService
RedisService.init_plugin_config()
plugin_config = RedisPluginConfig()
print(plugin_config)
result = plugin_config.get_plugin_config()
print(result)
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 vulnerability_details():
"""
获取漏洞详情
:return:
"""
try:
task_id = request.args.get("task_id")
current_user_name = session.get('user_name')
current_user = RedisService.get_user(current_user_name)
if int(current_user.role) < Role.ADMIN:
current_user_task_count = UserTaskService.count(
where=(UserTask.user_id == int(current_user.id),
UserTask.task_id == task_id))
if current_user_task_count == 0:
return jsonify(status=403,
message="查询失败",
data={"extra_info": "请勿尝试非法查看"})
return jsonify(
status=200,
message="查询成功",
vlun=VulnerabilityService.get_vulnerabilitys_nltdr(task_id))
except Exception as e:
logger.exception("vulnerability_details rasie error")
if isinstance(e, KeyError):
return jsonify(status=400,
message="查询失败",
data={"extra_info": "未传递参数task_id"})
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 addUser2dbAndRedis(self):
"""
同步一个账号到数据库和redis
:return:
"""
from api.service.redis_service import RedisService
from model.default_value import Role
from model.user import User, UserService
user_name = "b5mali4"
full_name = "小明"
email = "*****@*****.**"
dept_name = "信息安全部"
role = Role.USER
mobile_phone = "131xxxx9871"
if UserService.count(where=(User.user_name == user_name)) <= 0:
UserService.save(user_name=user_name,
full_name=full_name,
email=email,
dept_name=dept_name,
role=role,
mobile_phone=mobile_phone)
else:
UserService.update(fields=({
User.user_name: user_name,
User.full_name: full_name,
User.email: email,
User.dept_name: dept_name,
User.role: role,
User.mobile_phone: mobile_phone
}))
user = UserService.get_fields_by_where(
where=(User.user_name == user_name))[0]
user_info = {
"user_name": "b5mali4",
"full_name": "小明",
"email": "*****@*****.**",
"dept_name": "信息安全部",
"role": Role.USER,
"mobile_phone": "131xxxx9871",
"id": user.id
}
RedisService.update_user(user_name="b5mali4", user_info=user_info)
def modify_user(user_id):
"""
更新用户资料，主要包含权限，部门等等
data: {department: "信息安全部", emails: "", mobilephone: "18324742048", role: "管理员"}
:return:
"""
try:
put_data = request.get_json(force=True)
param_list = ["email", "mobile_phone", "role", "depart_ment"]
if has_dict_value_blank(put_data, param_list):
return jsonify(
status=400,
message="更新失败",
data={"extra_info": "请保证%s任一参数值不为空" % ','.join(param_list)})
email = put_data.get("email")
mobile_phone = put_data.get("mobile_phone")
role = put_data.get("role")
depart_ment = put_data.get("depart_ment")
UserService.update(fields=({
User.email: email,
User.mobile_phone: mobile_phone,
User.role: role,
User.dept_name: depart_ment
}),
where=(User.id == user_id))
user = UserService.get_fields_by_where(where=(User.id == user_id))[0]
RedisService.update_user(
user.user_name, {
"dept_name": depart_ment,
"role": role,
"mobile_phone": mobile_phone,
"email": email
})
return jsonify(status=200, message="更新用户成功", data={})
except Exception as e:
logger.exception("update_user error")
return jsonify(status=500,
message="未知异常",
data={"extra_info": "发生未知异常，请联系管理员查看异常日志"})
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 init_plugin_info():
"""
初始化插件信息到数据库中
:return:
"""
# 初始化redis配置信息
RedisService.init_plugin_config()
for checker_name, checker_instance in load_default_checkers().items():
if PluginInfoService.count(
where=(PluginInfo.plugin_name == checker_name)) == 0:
PluginInfoService.save(
author=checker_instance.info["author"],
plugin_name=checker_instance.info["name"],
plugin_tag=checker_instance.info["tag"],
imp_version=checker_instance.info["imp_version"],
description=checker_instance.info["description"],
repair=checker_instance.info["repair"],
type=checker_instance.info["type"]["fullname"],
chinese_type=checker_instance.info["type"]["fullchinesename"],
level=checker_instance.info["type"]["level"],
)
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 valid_role(current_user_name, role):
"""
判断当前用户权限是否符合API权限
如果用户登录过期，直接抛出异常，不同的登录认证方式自己可以做跳转之类的工作
:param current_user_name:
:param role:
:return:
"""
try:
user = RedisService.get_user(user_name=current_user_name)
return int(user.role) >= role
except Exception as e:
raise e
def auth_login(username, password, sessionid):
"""
代理认证模块，LDAP认证和账户认证，根据开关选择账号认证或者LDAP认证，如果两个开关都开启了，两种方式都支持
:param role:
:return:
"""
try:
user = RedisService.get_user(user_name=username)
if sessionid and hasattr(
user, "proxy_sessionid") and user.proxy_sessionid == sessionid:
return True, user
raise UserNotFoundInRedisException("账号%s proxy_sessionid不存在或者认证错误" %
username)
except UserNotFoundInRedisException:
if auth_login_logic(username, password):
# 设置proxy_session_id
proxy_sessionid = generate_session()
RedisService.update_user(username,
{"proxy_sessionid": proxy_sessionid})
user = RedisService.get_user(user_name=username)
return True, user
return False, None
def request_handler(self, request, user_info):
"""
将请求发送到MQ中
:param request:
:return:
Simple example code:
print(request.body_arguments)
print(request.headers)
print(request.body)
print(request.cookies)
print(request.version)
print(request.protocol)
print(request.host_name)
print(request.uri)
print(request.method)
"""
if not user_info:
return
task_id = user_info.current_task_id
current_user_name = user_info.user_name
raw_request_data = self.wrap_request(request, user_info)
# 是否为满足条件的请求
current_task = RedisService.get_task(task_id)
if current_task and "hook_rule" in current_task:
# *.xx.com
hook_rule = str(current_task.hook_rule).replace("*", ".*")
if not str(raw_request_data["data"]
["url"]).startswith(hook_rule) and re.match(
r'' + hook_rule, raw_request_data["data"]["url"],
re.S) is None:
return
if RedisService.create_urlclassifications(task_id, raw_request_data):
logger.info("满足正则条件,发送流量到MQ中")
scan_celery.delay(raw_request_data["data"], task_id,
current_user_name, TaskStatus.NONE)
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_vulnerability_count():
"""
获取当前用户所有扫描记录漏洞集合，按照type和level分类
:return:
"""
try:
current_user_name = session.get('user_name')
current_user = RedisService.get_user(current_user_name)
vlun_level_type_nums = VulnerabilityService.get_vulnerability_count(
current_user.id)
return jsonify(status=200, message="查询成功", data=vlun_level_type_nums)
except Exception as e:
logger.exception("vulnerability_count rasie error")
return jsonify(status=500,
message="未知异常",
data={"extra_info": "创建任务时出现未知异常，请联系管理员查看异常日志"})
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 get_user_info():
"""
获取用户信息，从session中获取用户名称，然后通过redis中查询，获取用户权限，根据用户权限来展示不同的页面，可以显示头像之类的
:return:
"""
try:
current_user_name = session.get("user_name")
current_user = RedisService.get_user(current_user_name)
return jsonify(status=200,
message="查询成功",
data={
"full_name": current_user.full_name,
"role": current_user.role
})
except Exception as e:
logger.exception("check_sso raise error")
return jsonify(status=500,
message="未知异常",
data={"extra_info": "发生未知异常，请联系管理员查看异常日志"})
def show_task_url_vulnerability_count():
"""
获取n天内的，任务，url，漏洞的数量
:return:
"""
try:
current_user_name = session.get('user_name')
current_user = RedisService.get_user(current_user_name)
day_range = int(request.args.get("day_range"))
tasks_urls_vulns_num = VulnerabilityService.get_tasks_urls_vulns_num_by_days(
user_id=current_user.id, day_range=day_range)
return jsonify(status=200, message="查询成功", data=tasks_urls_vulns_num)
except Exception as e:
logger.exception("show_task_url_vulnerability_count rasie error")
if isinstance(e, KeyError):
return jsonify(status=400,
message="查询失败",
data={"extra_info": "未传递参数day_range"})
return jsonify(status=500,
message="未知异常",
data={"extra_info": "创建任务时出现未知异常，请联系管理员查看异常日志"})
def retrieve_credentials(self):
"""
弹出账号密码基础认证，成功则写session
:return:
"""
auth_header = self.request.headers.get('Authorization', None)
proxy_session_id = self.get_cookie('proxy_sessionid', None)
if auth_header is not None:
# Basic Zm9vOmJhcg==
auth_mode, auth_base64 = auth_header.split(' ', 1)
assert auth_mode == 'Basic'
auth_username, auth_password = base64.b64decode(
auth_base64).decode("UTF-8").split(':', 1)
status, user_info = auth_login(auth_username, auth_password,
proxy_session_id)
# 认证失败
if not status:
self.write("认证失败,请确认账号密码是否正确")
self.set_status(401)
self.set_header('WWW-Authenticate', 'Basic realm="hunter"')
else:
self.set_cookie("proxy_sessionid",
user_info["proxy_sessionid"])
# 任务状态为关闭，或者任务不存在
if "current_task_id" not in user_info or (
"current_task_id" in user_info
and user_info["current_task_id"] != "" and
RedisService.get_task(user_info.current_task_id).status
!= str(TaskStatus.WORKING)):
self.write("后台无正在运行的任务,你需要重建一个新任务")
self.set_status(400)
self.finish()
status = False
return status, user_info
else:
self.set_status(401)
self.set_header('WWW-Authenticate', 'Basic realm="hunter"')
self.finish()
return False, None
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 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 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 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 insert_plugin():
"""
新增插件，
服务端:1.保存插件到本地，先上传到tmp目录，然后解析插件是否满足格式出tag并移动到tag目录，刷新插件列表 2.向MQ发送一条消息包含插件
引擎端:2.消费到MQ发送消息，并下载最新插件并新增到CHECKER_INSTANCE_DICT
:return:
"""
import shutil
from common.system_util import mkdir
from common.path import PLUGIN_PATH
from common.plugins_util import load_default_checkers
def allowed_file(filename):
"""
检测是否为python文件
:param filename:
:return:
"""
return '.' in filename and filename.rsplit('.', 1)[1] in ["py"]
def parse_plugin_file(plugin_file_path):
"""
warnning: 插件文件路径，这里会存在安全隐患，请按照实际需求考虑是否用imp.load_source
:param plugin_file_path:
:return:
"""
import imp
from exception.hunter_web_exception import BaseHunterException
# 解析插件并且分类
base_checker = imp.load_source('BaseChecker', plugin_file_path)
checker_instance = base_checker.Checker()
checker_instance.check_plugin_info()
# 检测插件是否重名
checker_name = checker_instance.info["name"]
if checker_name in load_default_checkers():
raise BaseHunterException("插件%s已经存在，请重新命名" % checker_name)
return checker_instance
def move_plugin(tmp_plugin_file_path, checker_instance, filename):
"""
:param checker_instance:
:return:
"""
# 移动插件到指定目录
tag = checker_instance.info["tag"]
mkdir(PLUGIN_PATH + tag)
formal_plugin_file_path = os.path.join(PLUGIN_PATH + tag, filename)
shutil.move(tmp_plugin_file_path, formal_plugin_file_path)
def save_plugin_info(checker_instance):
if PluginInfoService.count(where=(
PluginInfo.plugin_name == checker_instance.info["name"])) == 0:
PluginInfoService.save(
author=checker_instance.info["author"],
plugin_name=checker_instance.info["name"],
plugin_tag=checker_instance.info["tag"],
imp_version=checker_instance.info["imp_version"],
description=checker_instance.info["description"],
repair=checker_instance.info["repair"],
type=checker_instance.info["type"]["fullname"],
chinese_type=checker_instance.info["type"]["fullchinesename"],
level=checker_instance.info["type"]["level"],
)
try:
file = request.files['file']
if file and allowed_file(file.filename):
filename = file.filename
tmp_plugin_file_path = os.path.join(PLUGIN_PATH + "tmp/", filename)
file.save(tmp_plugin_file_path)
# 解析插件
checker_instance = parse_plugin_file(tmp_plugin_file_path)
move_plugin(tmp_plugin_file_path, checker_instance, filename)
# 保存到数据
save_plugin_info(checker_instance)
load_default_checkers(True)
RedisService.modify_plugin_switch(checker_instance, PluginSwith.ON)
system_notice_celery.delay(
broadcast={
"type": BroadCastType.PLUGIN,
"action": BroadCastAction.INSERT_PLUGIN,
"data": {
"name": checker_instance.info["name"]
}
})
return jsonify(status=200,
message="上传成功",
data={"extra_info": "您可以刷新网页后使用新插件"})
return jsonify(status=500,
message="上传失败",
data={"extra_info": "您上传的插件不是py文件"})
except Exception as e:
logger.exception("create_plugin raise error")
# 解析插件是否满足格式
return jsonify(status=500, message="上传失败", data={"extra_info": str(e)})
def testInitPluginInfo(self):
RedisService.init_plugin_info()
