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