90 lines
3.0 KiB
Python
90 lines
3.0 KiB
Python
#!/usr/bin/env python
|
|
# -*- coding: utf-8 -*-
|
|
# @Author : jeffzhang
|
|
# @Time : 18-5-14
|
|
# @File : plugin_management.py
|
|
# @Desc : ""
|
|
|
|
import time
|
|
import os
|
|
from flask import Flask, Blueprint, render_template, request, jsonify
|
|
from werkzeug.utils import secure_filename
|
|
from bson import ObjectId
|
|
from lib.mongo_db import connectiondb, db_name_conf
|
|
from fuxi.views.authenticate import login_check
|
|
from fuxi.views.modules.scanner.parse_plugin import parse_plugin
|
|
from instance import config
|
|
|
|
|
|
ProductionConfig = config.ProductionConfig
|
|
app = Flask(__name__)
|
|
app.config.from_object(ProductionConfig)
|
|
|
|
plugin_management = Blueprint('plugin_management', __name__)
|
|
tasks_db = db_name_conf()['tasks_db']
|
|
asset_db = db_name_conf()['asset_db']
|
|
server_db = db_name_conf()['server_db']
|
|
subdomain_db = db_name_conf()['subdomain_db']
|
|
vul_db = db_name_conf()['vul_db']
|
|
plugin_db = db_name_conf()['plugin_db']
|
|
|
|
|
|
# new plugin
|
|
@plugin_management.route('/new-asset', methods=['GET', 'POST'])
|
|
@login_check
|
|
def new_plugin():
|
|
pass
|
|
|
|
|
|
@plugin_management.route('/plugin-management', methods=['GET', 'POST'])
|
|
@login_check
|
|
def plugin_view():
|
|
# delete plugin
|
|
if request.method == "GET":
|
|
if request.args.get("delete"):
|
|
plugin_id = request.args.get('delete')
|
|
plugin_filename = connectiondb(plugin_db).find_one({"_id": ObjectId(plugin_id)})['plugin_filename']
|
|
if connectiondb(plugin_db).delete_one({'_id': ObjectId(plugin_id)}):
|
|
try:
|
|
os.remove(plugin_filename)
|
|
except Exception as e:
|
|
raise e
|
|
return "success"
|
|
else:
|
|
return "Warning"
|
|
# get plugin info
|
|
elif request.args.get("info"):
|
|
plugin_id = request.args.get('info')
|
|
plugin_info_data = connectiondb(plugin_db).find_one({'_id': ObjectId(plugin_id)})
|
|
del plugin_info_data['_id']
|
|
if plugin_info_data:
|
|
return jsonify(plugin_info_data)
|
|
else:
|
|
return jsonify({"result": "Warning"})
|
|
else:
|
|
# default view
|
|
plugin_info = connectiondb(plugin_db).find()
|
|
return render_template("plugin-management.html", plugin_info=plugin_info)
|
|
|
|
|
|
@plugin_management.route('/plugin-upload', methods=['GET', 'POST'])
|
|
@login_check
|
|
def plugin_upload():
|
|
file_path = app.config.get('POCSUITE_PATH')
|
|
file_data = request.files['file']
|
|
if file_data:
|
|
file_name = "_" + time.strftime("%y%m%d", time.localtime()) + "_" + secure_filename(file_data.filename)
|
|
save_path = file_path + file_name
|
|
file_data.save(save_path)
|
|
try:
|
|
new_plugin_info = parse_plugin(save_path)
|
|
if new_plugin_info:
|
|
db_insert = connectiondb(plugin_db).insert_one(new_plugin_info).inserted_id
|
|
if db_insert:
|
|
return jsonify({"result": "success"})
|
|
else:
|
|
return jsonify({"result": "Warning"})
|
|
except Exception as e:
|
|
print(e)
|
|
return "Warning"
|