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()