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 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 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_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 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": "未知异常,可以联系管理员到后台查看"})