def add_recommendations():
token = request.args.get('token')
check_token(token)
data = request.json
index = 0
while index < len(data):
AddRecommendation().execute(current_user, data[index])
index += 1
result = dict({"success": "yes"})
return jsonify(result)
def inner(*args, **kwargs):
tk = request.headers.get(
'Authorization'
) # this decorator should be called after [check_token]
uid = token.check_token(tk)
if not uid:
return msg.token_invalid_msg
else:
return func(*args, **kwargs)
def get(self):
tk = request.headers.get('Authorization')
uid = token.check_token(tk)
r = {'status': True, 'data': []}
task_list = md.Task.query.filter(md.Task.uid == uid).order_by(md.Task.id.desc()).all()
for task in task_list:
r['data'].append({'tid': task.id, 'tname': task.tname, 'date': task.date, 'status': task.status})
return r
def inner(*args, **kwargs):
tk = request.headers.get(
'Authorization'
) # this decorator should be called after [check_token]
uid = token.check_token(tk)
if not uid:
return msg.token_invalid_msg # login check as well
else:
user = md.User.query.filter_by(id=uid).first()
return func(*args, **kwargs) if user.is_admin else msg.deny_msg
def delete(self, task_id):
tk = request.headers.get('Authorization')
uid = token.check_token(tk)
task = md.Task.query.filter_by(id=task_id, uid=uid).first()
if task:
db.session.delete(task)
db.session.commit()
r = msg.success_msg
else:
r = msg.no_resource_msg
return r
def put(self, task_id):
tk = request.headers.get('Authorization')
uid = token.check_token(tk)
task = md.Task.query.filter_by(id=task_id, uid=uid).first()
user = md.User.query.filter_by(id=uid).first()
if task:
# to create a new process to start distributed task, making this request immediately return
scan_process = multiprocessing.Process(target=celery_task.launch, args=(task.id, tk, user.email))
scan_process.start()
r = msg.success_msg
else:
r = msg.no_resource_msg
return r
def post(self):
args, tk = self.parser.parse_args(), request.headers.get('Authorization')
uid = token.check_token(tk)
date = time.strftime("%Y-%m-%d %H:%M", time.localtime())
new_task = md.Task(uid=uid, tname=args['tname'], date=date, status=0)
db.session.add(new_task)
db.session.commit()
host_list = json.loads(args['hosts']) # convert pure json to python list
for host in host_list:
if dns.valid_host(host['target']):
new_host = md.Host(tid=new_task.id, hname=new_task.tname, target=host['target'], policy=host['policy'])
db.session.add(new_host)
db.session.commit()
return msg.success_msg
def __call__(self, request):
request_url = request.path_info
for p in white_list:
if request_url in p:
response = self.get_response(request)
return response
for p in black_list:
if request_url in p:
response = HttpResponse()
response.content= json.dumps({"code":"9999","message":"请求内容非法","data":None})
response["Content-Type"]="application/json;charset=UTF-8"
return response
token = request.META.get("HTTP_TOKEN") # 获取请求头中token的值
if check_token(token):
response = self.get_response(request)
return response
response = HttpResponse()
response.content = json.dumps({"code": "9999", "message": "暂未登录或token已过期", "data": None})
response["Content-Type"] = "application/json"
return response