def TaskCount():
result = {}
result['running'] = str(Status.objects(status=STATUS.RUN).count())
result['waiting'] = str(Status.objects(status=STATUS.WAIT).count())
result['finish'] = str(Status.objects(status=STATUS.FINISH).count())
result['fail'] = str(Status.objects(status=STATUS.FAIL).count())
return json.dumps(result)
def TaskList(show_type):
result = []
page_result = {}
if show_type == 'index':
lens = len(Status.objects)
if lens > 10:
data = Status.objects[lens - 10:lens]
else:
data = Status.objects().all()
elif show_type == 'all':
data = Status.objects().all()
elif show_type == 'page':
page_number = int(request.args.get('p'))
lens = len(Status.objects)
page_result['total_page'] = str(lens//page_item_number+1)
if page_number > lens//page_item_number+1:
return json.dumps(page_result)
if page_number * page_item_number < lens:
data = Status.objects[lens-(page_number*page_item_number):lens-(page_item_number*(page_number-1))]
else:
data = Status.objects[0 :lens-(page_item_number*(page_number-1))]
for each in data:
info = {}
info['create_time'] = str(each.create_time).split('.')[0]
info['task_name'] = cgi.escape(each.task_name)
info['task_id'] = each.task_id
info['status'] = each.status
info['progress'] = each.progress
info['warning'] = each.warning
result.append(info)
result.reverse()
if show_type == 'page':
page_result['info'] = result
return json.dumps(page_result)
return json.dumps(result)
def AddTaskToDataBase(task_id, task_name):
if not Status.objects(task_id=task_id).count():
Status(task_id=task_id,
task_name=task_name,
warning='',
progress='0',
status='Waiting').save()
def SaveWarningToDatabase(data):
if args.use_database == False:
return
try:
status = Status.objects(task_id=args.task_id)
status.update(warning=data)
except Exception, e:
raise DatabaseException(repr(e))
def TaskInfo(task_id):
result = {}
result['info'] = {}
status = Status.objects(task_id=task_id).first()
if not status:
result['success'] = False
result['message'] = u'任务ID不存在。'
return json.dumps(result)
if status.status == STATUS.WAIT:
result['success'] = False
result['message'] = u'任务尚未开始，请等待。'
return json.dumps(result)
if status.status == STATUS.RUN:
result['success'] = False
result['message'] = u'扫描正在进行中，请等待。'
return json.dumps(result)
if status.status == STATUS.CLOSE:
result['success'] = False
result['message'] = u'该任务已关闭，无法查看。'
return json.dumps(result)
if status.status == STATUS.FAIL:
result['success'] = False
result['message'] = str(status.warning)
return json.dumps(result)
task = Task.objects(task_id=task_id).first()
if not task:
result['success'] = False
result['message'] = u'任务ID不存在.'
return json.dumps(result)
scan_result = Result.objects(task_id=task_id).first()
result['info']['task_name'] = cgi.escape(task.task_name)
result['info']['input_target'] = task.short_target
result['info']['scan_mode'] = task.scan_mode
result['info']['target_type'] = task.target_type
result['info']['script_type'] = task.script_type
result['info']['start_time'] = str(task.start_time).split('.')[0]
result['info']['finish_time'] = str(task.finish_time).split('.')[0]
result['info']['target_number'] = task.target_number
result['info']['full_target'] = task.full_target
result['info']['script_number'] = len(json.loads(task.script_info))
result['info']['high_count'] = scan_result.high_count
result['info']['medium_count'] = scan_result.medium_count
result['info']['low_count'] = scan_result.low_count
result['info']['result'] = scan_result.result
scripts_list = []
for script_name in json.loads(task.script_info):
script_data = {}
script = Script.objects(script_name=script_name).first()
script_data['script_name'] = script_name
script_data['script_title'] = script.script_title
script_data['script_id'] = script.script_id
script_data['script_level'] = script.script_level
scripts_list.append(script_data)
result['info']['scripts_info'] = scripts_list
result['info']['warning'] = status.warning
result['success'] = True
result['message'] = ''
return json.dumps(result)
def UpdateStatus():
if args.use_database == False:
return
try:
if Status.objects(task_id=args.task_id).count():
status = Status.objects(task_id=args.task_id).first()
status.update(status=STATUS.RUN)
else:
try:
Status(task_id=args.task_id,
task_name=args.task_name,
warning='',
progress='0',
status=STATUS.RUN).save()
except Exception, e:
if 'duplicate key' in e:
pass
except Exception, e:
raise DatabaseException(repr(e))
def get_last_scape_dt(mobile):
conn = connect_mongodb()
last_status = Status.objects(mobile=mobile).order_by("-scrape_dt").first()
conn.close()
if last_status:
last_dt = last_status.scrape_dt
# 返回样例 (u'2016-11-04 17:15:04', u'2016-11', u'201611')
return last_dt, last_dt[:7], last_dt[:7].replace("-", ""), last_dt[:10] + ' 00:00:00'
else:
return None, None, None, None
def DeleteTask(task_id):
result = {}
if Status.objects(task_id=task_id).count():
status = Status.objects(task_id=task_id).first()
if status.status == STATUS.WAIT or status.status == STATUS.RUN:
try:
celery.control.revoke(task_id, terminate=True)
status.update(status=STATUS.CLOSE)
result['success'] = True
except:
result['success'] = False
else:
try:
status.delete()
task = Task.objects(task_id=task_id).first()
if task:
task.delete()
result['success'] = True
except:
result['success'] = False
else:
result['success'] = False
return json.dumps(result)
def DataSearch(search_type):
keyword = request.args.get('keyword')
page_number = int(request.args.get('p'))
result = []
page_result = {}
if search_type == 'script':
lens = len(
Script.objects(__raw__={
'script_info': re.compile(keyword)
}).all())
page_result['total_page'] = str(lens // page_item_number + 1)
if page_number > lens // page_item_number + 1:
return json.dumps(page_result)
if page_number * page_item_number < lens:
scripts = Script.objects(
__raw__={'script_info': re.compile(keyword)
})[lens - (page_number * page_item_number):lens -
(page_item_number * (page_number - 1))]
else:
scripts = Script.objects(
__raw__={'script_info': re.compile(keyword)
})[0:lens - (page_item_number * (page_number - 1))]
for each in scripts:
info = {}
info['id'] = each.script_id
info['detail'] = each.script_info
info['create_time'] = each.script_update_time
info['author'] = each.script_author
info['level'] = each.script_level
info['count'] = Vuln.objects(script=each.script_id).count()
info['title'] = each.script_title
result.append(info)
elif search_type == 'task':
lens = len(
Status.objects(__raw__={
'task_name': re.compile(keyword)
}).all())
page_result['total_page'] = str(lens // page_item_number + 1)
if page_number > lens // page_item_number + 1:
return json.dumps(page_result)
if page_number * page_item_number < lens:
tasks = Status.objects(__raw__={'task_name': re.compile(
keyword)})[lens - (page_number * page_item_number):lens -
(page_item_number * (page_number - 1))]
else:
tasks = Status.objects(__raw__={'task_name': re.compile(
keyword)})[0:lens - (page_item_number * (page_number - 1))]
for each in tasks:
info = {}
info['create_time'] = str(each.create_time).split('.')[0]
info['task_name'] = cgi.escape(each.task_name)
info['task_id'] = each.task_id
info['status'] = each.status
info['progress'] = each.progress
result.append(info)
elif search_type == 'vuln':
lens = len(Vuln.objects(__raw__={'target': re.compile(keyword)}).all())
page_result['total_page'] = str(lens // page_item_number + 1)
if page_number > lens // page_item_number + 1:
return json.dumps(page_result)
if page_number * page_item_number < lens:
vulns = Vuln.objects(__raw__={'target': re.compile(
keyword)})[lens - (page_number * page_item_number):lens -
(page_item_number * (page_number - 1))]
else:
vulns = Vuln.objects(
__raw__={'target': re.compile(keyword)})[0:lens -
(page_item_number *
(page_number - 1))]
for each in vulns:
info = {}
script = Script.objects(script_id=each.script).first()
try:
info['task_name'] = cgi.escape(
Task.objects(task_id=each.task_id).first().task_name)
except:
info['task_name'] = u'该任务已删除，无法查看'
info['task_id'] = each.task_id
info['target'] = each.target
info['script_name'] = script.script_name
info['script_id'] = each.script
info['message'] = each.message
info['script_type'] = each.script_type
info['time'] = str(each.create_time).split('.')[0]
info['level'] = script.script_level
result.append(info)
page_result['info'] = result
return json.dumps(page_result)