完成规则管理后台部分、以及GitTool类和decompress类。

app/__init__.py

@@ -1,16 +1,16 @@
 # coding: utf-8
-import sys
-import os
 import ConfigParser
+import os
+import sys
 
 from flask import Flask
+from flask.ext.migrate import MigrateCommand, Migrate
 from flask.ext.script import Manager, Server
 from flask.ext.sqlalchemy import SQLAlchemy
-from flask.ext.migrate import MigrateCommand, Migrate
+from flask.ext.bootstrap import Bootstrap
 
 from utils import log
 
-
 log.info('Initialization HTTP Server')
 reload(sys)
 sys.setdefaultencoding('utf-8')
@@ -24,6 +24,8 @@
 web.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = True
 web.config['SQLALCHEMY_DATABASE_URI'] = config.get('database', 'mysql')
 
+bootstrap = Bootstrap(web)
+
 db = SQLAlchemy(web)
 
 # just use the migration script's app context when you import the models
@@ -41,7 +43,9 @@
 manager.add_command('db', MigrateCommand)
 manager.add_command('runserver', Server(host=host, port=port))
 
-from app import route
+from app.controller import route
+from app.controller import RulesAdmin
+
 log.info('Cobra HTTP Server Started')
 
 

app/controller/RulesAdmin.py

@@ -0,0 +1,236 @@
+#!/usr/bin/env python
+import time
+
+from flask import render_template, request, jsonify
+
+from app import web, CobraRules, CobraVuls, db, CobraSupportLanguage
+
+# default admin url
+ADMIN_URL = '/admin'
+
+
+@web.route(ADMIN_URL + '/', methods=['GET'])
+@web.route(ADMIN_URL + '/index', methods=['GET'])
+def index():
+    return 'admin/index - todo: login page'
+
+
+# main view
+@web.route(ADMIN_URL + '/main', methods=['GET'])
+def main():
+    return render_template("rulesadmin/main.html")
+
+
+# all rules button
+@web.route(ADMIN_URL + '/rules', methods=['GET'])
+def rules():
+    # cobra_rules = CobraRules.query.paginate(1, per_page=5, error_out=False)
+    cobra_rules = CobraRules.query.all()
+    cobra_vuls = CobraVuls.query.all()
+    cobra_lang = CobraSupportLanguage.query.all()
+    all_vuls = {}
+    all_language = {}
+    for vul in cobra_vuls:
+        all_vuls[vul.id] = vul.name
+    for lang in cobra_lang:
+        all_language[lang.id] = lang.language
+
+    # replace id with real name
+    for rule in cobra_rules:
+        try:
+            rule.vul_id = all_vuls[rule.vul_id]
+        except KeyError:
+            rule.vul_id = 'Unknown Type'
+        try:
+            rule.language = all_language[rule.language]
+        except KeyError:
+            rule.language = 'Unknown Language'
+
+    data = {
+        # 'paginate': cobra_rules,
+        'rules': cobra_rules,
+    }
+
+    return render_template('rulesadmin/rules.html', data=data)
+
+
+# add new rules button
+@web.route(ADMIN_URL + '/add_new_rule', methods=['GET', 'POST'])
+def add_new_rule():
+    if request.method == 'POST':
+        vul_type = request.form['vul_type']
+        lang = request.form['language']
+        regex = request.form['regex']
+        description = request.form['description']
+
+        if not vul_type or vul_type == "":
+            return jsonify(tag='danger', msg='vul type error.')
+        if not lang or lang == "":
+            return jsonify(tag='danger', msg='language error.')
+        if not regex or regex == "":
+            return jsonify(tag='danger', msg='regex can not be blank')
+        if not description or description == "":
+            return jsonify(tag='danger', msg='description can not be blank')
+
+        current_time = time.strftime('%Y-%m-%d %X', time.localtime())
+        rule = CobraRules(vul_type, lang, regex, description, current_time, current_time)
+        try:
+            db.session.add(rule)
+            db.session.commit()
+            return jsonify(tag='success', msg='add success.')
+        except:
+            return jsonify(tag='danger', msg='add failed, try again later?')
+    else:
+        vul_type = CobraVuls.query.all()
+        languages = CobraSupportLanguage.query.all()
+        data = {
+            'vul_type': vul_type,
+            'languages': languages
+        }
+        return render_template('rulesadmin/add_new_rule.html', data=data)
+
+
+# del special rule
+@web.route(ADMIN_URL + '/del_rule', methods=['POST'])
+def del_rule():
+    vul_id = request.form['rule_id']
+    if vul_id:
+        r = CobraRules.query.filter_by(id=vul_id).first()
+        try:
+            db.session.delete(r)
+            db.session.commit()
+            return jsonify(tag='success', msg='delete success.')
+        except:
+            return jsonify(tag='danger', msg='delete failed. Try again later?')
+    else:
+        return jsonify(tag='danger', msg='wrong id')
+
+
+# edit special rule
+@web.route(ADMIN_URL + '/edit_rule/<int:rule_id>', methods=['GET', 'POST'])
+def edit_rule(rule_id):
+    if request.method == 'POST':
+        vul_type = request.form['vul_type']
+        lang = request.form['language']
+        regex = request.form['regex']
+        description = request.form['description']
+        rule_id = request.form['rule_id']
+
+        if not vul_type or vul_type == "":
+            return jsonify(tag='danger', msg='vul type error.')
+        if not lang or lang == "":
+            return jsonify(tag='danger', msg='language error.')
+        if not regex or regex == "":
+            return jsonify(tag='danger', msg='regex can not be blank')
+        if not description or description == "":
+            return jsonify(tag='danger', msg='description can not be blank')
+
+        r = CobraRules.query.filter_by(id=rule_id).first()
+        r.vul_id = vul_type
+        r.language = lang
+        r.regex = regex
+        r.description = description
+        try:
+            db.session.add(r)
+            db.session.commit()
+            return jsonify(tag='success', msg='save success.')
+        except:
+            return jsonify(tag='danger', msg='save failed. Try again later?')
+    else:
+        r = CobraRules.query.filter_by(id=rule_id).first()
+        vul_type = CobraVuls.query.all()
+        languages = CobraSupportLanguage.query.all()
+        return render_template('rulesadmin/edit_rule.html', data={
+            'vul_type': r.vul_id,
+            'language': r.language,
+            'regex': r.regex,
+            'description': r.description,
+            'all_vuls': vul_type,
+            'all_lang': languages,
+        })
+
+
+# add new vuls button
+@web.route(ADMIN_URL + '/add_new_vul', methods=['GET', 'POST'])
+def add_new_vul():
+    if request.method == 'POST':
+        name = request.form['name']
+        description = request.form['description']
+        if not name or name == "":
+            return jsonify(tag='danger', msg='name is empty')
+        if not description or description == "":
+            return jsonify(tag='danger', msg='description is empty')
+
+        current_time = time.strftime('%Y-%m-%d %X', time.localtime())
+        vul = CobraVuls(name, description, current_time, current_time)
+        try:
+            db.session.add(vul)
+            db.session.commit()
+            return jsonify(tag='success', msg='Add Success.')
+        except:
+            return jsonify(tag='danger', msg='Add failed. Please try again later.')
+
+    else:
+        return render_template('rulesadmin/add_new_vul.html')
+
+
+# show all vuls click
+@web.route(ADMIN_URL + '/vuls', methods=['GET'])
+def vuls():
+    all_vuls = CobraVuls.query.all()
+    data = {
+        'vuls': all_vuls
+    }
+    return render_template('rulesadmin/vuls.html', data=data)
+
+
+# del special vul
+@web.route(ADMIN_URL + '/del_vul', methods=['POST'])
+def del_vul():
+    vul_id = request.form['vul_id']
+    if vul_id:
+        v = CobraVuls.query.filter_by(id=vul_id).first()
+        try:
+            db.session.delete(v)
+            db.session.commit()
+            return jsonify(tag='success', msg='delete success.')
+        except:
+            return jsonify(tag='danger', msg='delete failed. Try again later?')
+    else:
+        return jsonify(tag='danger', msg='wrong id')
+
+
+# edit special vul
+@web.route(ADMIN_URL + '/edit_vul/<int:vul_id>', methods=['GET', 'POST'])
+def edit_vul(vul_id):
+    if request.method == 'POST':
+        name = request.form['name']
+        description = request.form['description']
+        if not name or name == "":
+            return jsonify(tag='danger', msg='name can not be empty')
+        if not description or description == "":
+            return jsonify(tag='danger', msg='description can not be empty')
+        v = CobraVuls.query.filter_by(id=vul_id).first()
+        v.name = name
+        v.description = description
+        try:
+            db.session.add(v)
+            db.session.commit()
+            return jsonify(tag='success', msg='save success.')
+        except:
+            return jsonify(tag='danger', msg='save failed. Try again later?')
+    else:
+        v = CobraVuls.query.filter_by(id=vul_id).first()
+        return render_template('rulesadmin/edit_vul.html', data={
+            'name': v.name,
+            'description': v.description,
+        })
+
+
+# api: get all rules count
+@web.route(ADMIN_URL + '/all_rules_count', methods=['GET'])
+def all_rules_count():
+    rules_count = CobraRules.query.count()
+    return str(rules_count)
+
+

app/route.py → app/controller/route.py

@@ -14,10 +14,12 @@
 import os
 import time
 import argparse
+import ConfigParser
 
 import magic
 from utils import log
 from flask import request, jsonify, render_template
+from werkzeug import secure_filename
 
 from app import web, CobraTaskInfo, db
 
@@ -72,23 +74,30 @@ def add():
         # no files, should check username and password
         task_type = 1
         url = request.form['url']
-        username = request.form['username']
-        password = request.form['password']
+        username = request.form['username'] if request.form['username'] != '' else None
+        password = request.form['password'] if request.form['password'] != '' else None
+        branch = request.form['branch'] if request.form['branch'] != '' else 'master'
 
-        if not url or not username or not password:
-            return jsonify(code=1002, msg=u'please support username, password and gitlab.')
+        if not url:
+            return jsonify(code=1002, msg=u'please support gitlab url. '
+                                          u'If this is a public repo, just leave username and password blank')
 
         # insert into db
-        new_task = CobraTaskInfo(task_type, int(time.time()), None, url, username, password, scan_type, level,
+        new_task = CobraTaskInfo(task_type, int(time.time()), None, url, branch, username, password, scan_type, level,
                                  scan_way, old_version, new_version)
         db.session.add(new_task)
         db.session.commit()
     else:
         # there is a file, check file format and uncompress it.
+        # get uploads directory
+        config = ConfigParser.ConfigParser()
+        config.read('config')
+        upload_directory = config.get('cobra', 'upload_directory') + os.sep
+
         task_type = 2
         upload_src = request.files['file']
-        filename = str(int(time.time())) + '_' + upload_src.filename
-        filepath = 'uploads/' + filename
+        filename = str(int(time.time())) + '_' + secure_filename(upload_src.filename)
+        filepath = upload_directory + filename
         upload_src.save(filepath)
 
         # if you upload a rar file, upload_src.mimetype will returns "application/octet-stream"
@@ -99,7 +108,7 @@ def add():
             os.remove(filepath)
             return jsonify(code=1002, msg=u'only rar, zip and tar.gz supported.')
 
-        new_task = CobraTaskInfo(task_type, int(time.time()), filename, None, None, None, scan_type, level,
+        new_task = CobraTaskInfo(task_type, int(time.time()), filename, None, None, None, None, scan_type, level,
                                  scan_way, old_version, new_version)
         db.session.add(new_task)
         db.session.commit()