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 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 authorize_route():
"""
ldap认证账号体系
:return:
"""
try:
post_data = request.get_json(force=True)
post_user_name = post_data.get("user_name")
post_pass_word = post_data.get("pass_word")
ldap_config = LdapConfigService.get_single_instance()
if ldap_config.ldap_switch is False:
return jsonify(status=500,
message="登录失败",
data={"extra_info": "不支持ldap认证，请后台配置并开启ldap模块"})
status, result_dict = ldap_auth(post_user_name, post_pass_word)
if status:
user_name = result_dict["user_name"]
# 保存更新数据库和Redis
if UserService.count(where=(User.user_name == user_name)) <= 0:
UserService.save(user_name=result_dict["user_name"],
full_name=result_dict["full_name"],
dept_name=result_dict["dept_name"],
email=result_dict["email"],
mobile_phone=result_dict["mobile"])
else:
UserService.update(fields=({
User.full_name:
result_dict["full_name"],
User.dept_name:
result_dict["dept_name"],
User.email:
result_dict["email"],
User.mobile_phone:
result_dict["mobile"]
}),
where=(User.user_name == user_name))
db_user = UserService.get_fields_by_where(
where=(User.user_name == user_name))[0]
BaseAuthModule.modify_user_info_cache_session(
user_name=db_user.user_name, db_user=db_user)
return jsonify(status=200,
message="认证成功",
data={
"extra_info":
"稍后自动跳转首页,请耐心等待",
"site":
get_system_config()['front_end']['index']
})
return jsonify(status=403,
message="认证出错",
data={
"extra_info": "账号密码登录出错",
"site": "/login"
})
except Exception as e:
logger.exception("auth_account raise 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": "创建任务时出现未知异常，请联系管理员查看异常日志"})