Merge pull request #382 from InQuest/issue/381-ability-to-revert-prev…

…ious-yara

#381 : Adding ability to revert to previous Yara rule from the UI.

app/routes/yara_rules.py

@@ -1,4 +1,9 @@
+import re
+
+from sqlalchemy import and_
+
 from app import app, db, auto, ENTITY_MAPPING
+from app.models.cfg_category_range_mapping import CfgCategoryRangeMapping
 from app.models.yara_rule import Yara_rule_history, Yara_rule
 from app.routes import test_yara_rule
 from app.models import yara_rule, cfg_states, comments
@@ -98,7 +103,6 @@ def get_all_yara_rules():
     include_tags = bool(distutils.util.strtobool(request.args.get('include_tags', "true")))
     include_comments = bool(distutils.util.strtobool(request.args.get('include_comments', "true")))
 
-
     if include_yara_string:
         include_yara_string = True
 
@@ -306,6 +310,7 @@ def activate_yara_rule(id):
     db.session.commit()
     return jsonify(entity.to_dict()), 201
 
+
 @app.route('/ThreatKB/yara_rules/<int:id>', methods=['PUT'])
 @auto.doc()
 @login_required
@@ -535,7 +540,7 @@ def copy_yara_rules():
     Return: yara strings for copy"""
 
     signatures = []
-    if 'copy' in request.json and request.json['copy']\
+    if 'copy' in request.json and request.json['copy'] \
             and 'ids' in request.json['copy'] and request.json['copy']['ids']:
         for sig_id in request.json['copy']['ids']:
             sig = yara_rule.Yara_rule.query.get(sig_id)
@@ -562,3 +567,99 @@ def delete_all_inactive_yara_rules():
     db.session.query(yara_rule.Yara_rule).filter(yara_rule.Yara_rule.active == 0).delete()
     db.session.commit()
     return jsonify(''), 200
+
+
+@app.route('/ThreatKB/yara_rules/<int:yara_rule_id>/revert-to-revision/<int:revision>', methods=['PUT'])
+@auto.doc()
+@login_required
+def revert_yara_rule_to_revision(yara_rule_id, revision):
+    """Revert given yara rule to provided revision number
+    Return: Success Code"""
+
+    current_entity = yara_rule.Yara_rule.query.get(yara_rule_id)
+    revision_entity = Yara_rule_history.query \
+        .filter_by(yara_rule_id=yara_rule_id) \
+        .filter_by(revision=revision) \
+        .first()
+
+    if not current_entity or not revision_entity:
+        abort(404)
+
+    if not current_user.admin and current_entity.owner_user_id != current_user.id:
+        abort(403)
+
+    revision_dict = revision_entity.to_dict()
+    yara_revision_dict = revision_dict["rule_json"]
+
+    temp_sig_id = current_entity.eventid
+    if not current_entity.category == yara_revision_dict['category']:
+        temp_sig_id = CfgCategoryRangeMapping.get_next_category_eventid(yara_revision_dict['category'])
+
+    db.session.add(yara_rule.Yara_rule_history(date_created=datetime.datetime.now(), revision=current_entity.revision,
+                                               rule_json=json.dumps(current_entity.to_revision_dict()),
+                                               user_id=current_user.id,
+                                               yara_rule_id=current_entity.id,
+                                               state=current_entity.state))
+
+    current_entity = yara_rule.Yara_rule(
+        state=yara_revision_dict['state'],
+        name=yara_revision_dict['name'],
+        description=yara_revision_dict['description'],
+        references=yara_revision_dict['references'],
+        category=yara_revision_dict['category'],
+        condition=re.sub('^condition:\\n\\t', '', yara_revision_dict['condition']),
+        strings=re.sub('^strings:\\n\\t', '', yara_revision_dict['strings']),
+        eventid=temp_sig_id,
+        id=yara_rule_id,
+        creation_date=yara_revision_dict['creation_date'],
+        modified_user_id=current_user.id,
+        last_revision_date=datetime.datetime.now(),
+        owner_user_id=yara_revision_dict['owner_user']['id'] if yara_revision_dict['owner_user'] else None,
+        revision=current_entity.revision + 1,
+        imports=yara_revision_dict['imports'],
+        active=yara_revision_dict['active'],
+        mitre_techniques=yara_revision_dict['mitre_techniques'],
+        mitre_tactics=yara_revision_dict['mitre_tactics'],
+        files=yara_revision_dict['files'],
+    )
+
+    mitre_techniques = Cfg_settings.get_setting("MITRE_TECHNIQUES").split(",")
+    matches = [technique for technique in current_entity.mitre_techniques if technique not in mitre_techniques]
+    if matches:
+        raise (Exception(
+            "The following techniques were not found in the configuration: %s. Check 'MITRE_TECHNIQUES' on the settings page" % (
+                matches)))
+
+    mitre_tactics = Cfg_settings.get_setting("MITRE_TACTICS").split(",")
+    matches = [tactic for tactic in current_entity.mitre_tactics if tactic not in mitre_tactics]
+    if matches:
+        raise (Exception(
+            "The following tactics were not found in the configuration: %s. Check 'MITRE_TACTICS' on the settings page" % (
+                matches)))
+
+    db.session.merge(current_entity)
+    db.session.commit()
+
+    dirty = False
+    for name, value_dict in yara_revision_dict['metadata_values'].iteritems():
+        if not name or not value_dict:
+            continue
+
+        m = db.session.query(MetadataMapping).join(Metadata, Metadata.id == MetadataMapping.metadata_id).filter(
+            Metadata.key == name).filter(Metadata.artifact_type == ENTITY_MAPPING["SIGNATURE"]).filter(
+            MetadataMapping.artifact_id == current_entity.id).first()
+        if m:
+            m.value = value_dict["value"]
+            db.session.add(m)
+            dirty = True
+        else:
+            m = db.session.query(Metadata).filter(Metadata.key == name).filter(
+                Metadata.artifact_type == ENTITY_MAPPING["SIGNATURE"]).first()
+            db.session.add(MetadataMapping(value=value_dict["value"], metadata_id=m.id, artifact_id=current_entity.id,
+                                           created_user_id=current_user.id))
+            dirty = True
+
+    if dirty:
+        db.session.commit()
+
+    return jsonify(''), 204

app/static/js/yara_rule/yara_rule-controller.js

@@ -459,6 +459,9 @@ angular.module('ThreatKB')
                         }, function (error) {
                             growl.error(error, {ttl: -1})
                         })
+                    } else if (id_or_rule.reverted) {
+                        growl.info("Successfully reverted '" + $scope.yara_rule.name + "' to revision '" + id_or_rule.revertedToRevision + "'", {ttl: 3000});
+                        getPage();
                     } else {
                         id = id_or_rule.id;
                         $scope.yara_rule = id_or_rule;
@@ -575,7 +578,7 @@ angular.module('ThreatKB')
                 });
 
                 yara_ruleSave.result.then(function (entity) {
-                    if (entity.merge) {
+                    if (entity.merge || entity.reverted) {
                         $scope.save(entity);
                     } else {
                         $scope.yara_rule = entity;
@@ -719,7 +722,7 @@ angular.module('ThreatKB')
             };
 
             $scope.save_artifact = function () {
-
+                $scope.yara_rule.do_not_bump_revision = $scope.do_not_bump_revision;
                 Yara_rule.resource.update({id: $scope.yara_rule.id}, $scope.yara_rule,
                     function (data) {
                         if (!data) {
@@ -821,7 +824,13 @@ angular.module('ThreatKB')
 
             $scope.cfg_states = Cfg_states.query();
             $scope.cfg_category_range_mapping = CfgCategoryRangeMapping.query();
-            $scope.do_not_bump_revision = true;
+            if ($scope.do_not_bump_revision == null) {
+                $scope.do_not_bump_revision = true;
+            }
+
+            $scope.toggle_bump_release = function () {
+                $scope.do_not_bump_revision = this.do_not_bump_revision;
+            };
 
             $scope.just_opened = true;
             $scope.negTestDir = Cfg_settings.get({key: "NEGATIVE_TESTING_FILE_DIRECTORY"});
@@ -900,6 +909,7 @@ angular.module('ThreatKB')
             };
 
             $scope.ok = function () {
+                $scope.yara_rule.do_not_bump_revision = $scope.do_not_bump_revision;
                 // Check if outstanding comment
                 if ($scope.yara_rule.new_comment && $scope.yara_rule.new_comment.trim() && confirm('There is a unsaved comment, do you wish to save it as well?')) {
                     $scope.add_comment($scope.yara_rule.id);
@@ -986,6 +996,18 @@ angular.module('ThreatKB')
                 $uibModalInstance.close($scope.selected_signature);
             };
 
+            $scope.revertRevision = function (id, revision) {
+                Yara_rule.revertRevision(id, revision)
+                    .then(function (response) {
+                        $scope.yara_rule = Yara_rule.resource.get({id: id, include_yara_string: 1});
+                        $scope.yara_rule.reverted = true;
+                        $scope.yara_rule.revertedToRevision = revision;
+                        $uibModalInstance.close($scope.yara_rule);
+                    }, function (error) {
+                        growl.error(error.data, {ttl: -1});
+                    });
+            };
+
         }])
     .controller('Yara_ruleViewController', ['$scope', '$uibModalInstance', 'yara_rule', '$location', '$window', '$cookies',
         function ($scope, $uibModalInstance, yara_rule, $location, $window, $cookies) {
@@ -1071,7 +1093,7 @@ angular.module('ThreatKB')
         }])
     .controller('Yara_revisionController', ['$scope', 'Yara_rule',
         function ($scope, Yara_rule) {
-            $scope.revision_diff = null
+            $scope.revision_diff = null;
             $scope.calculateRevisionDiff = function () {
                 if (!$scope.selectedRevisions.compared) {
                     $scope.revision_diff = null;

app/static/js/yara_rule/yara_rule-service.js

@@ -91,6 +91,17 @@ angular.module('ThreatKB')
                 }
             );
         }
+        function revertRevision(sig_id, revision) {
+            return $http.put('/ThreatKB/yara_rules/' + sig_id + '/revert-to-revision/' + revision)
+                .then(function (success) {
+                    if (success.status === 200) {
+                        return success.data;
+                    }
+                }, function (error) {
+                    return $q.reject(error.data);
+                }
+            );
+        }
         return {
             resource: $resource('ThreatKB/yara_rules/:id', {}, {
                 'query': {method: 'GET', isArray: true},
@@ -102,6 +113,7 @@ angular.module('ThreatKB')
             updateBatch: updateBatch,
             deleteBatch: deleteBatch,
             getSignatureFromRevision: getSignatureFromRevision,
+            revertRevision: revertRevision,
             activateRule: activateRule,
             delete_all_inactive: delete_all_inactive
         };

app/static/views/yara_rule/yara_rules.html

@@ -603,6 +603,14 @@ <h4 class="modal-title" id="myC2ipLabel" style="float: left">
                                     </option>
                                 </select>
                             </div>
+                            <div style="max-width: min-content">
+                                <label>&nbsp;</label>
+                                <button type="button" class="btn btn-danger"
+                                        ng-disabled="!selectedRevisions.compared"
+                                        confirmed-click="revertRevision(yara_rule.id, selectedRevisions.compared.revision)"
+                                        ng-confirm-click="Are you sure you want to revert to revision?">Revert
+                                </button>
+                            </div>
                         </div>
                         <div class="yara-revisions-view" ng-controller="Yara_revisionController">
                                 <pre ng-if="revision_diff"><code ng-bind-html="revision_diff"></code></pre>
@@ -639,7 +647,7 @@ <h4 class="modal-title" id="myC2ipLabel" style="float: left">
                     </button>
                 </div>
                 <div class="pull-right" style="padding: 5px 10px 0 5px;" ng-if="yara_rule.id">
-                    <input type="checkbox" class="pull-left" ng-model="do_not_bump_revision">
+                    <input type="checkbox" class="pull-left" ng-model="do_not_bump_revision" ng-change="toggle_bump_release()">
                     &nbsp;Do not bump revision
                 </div>
             </div>

requirements.txt

@@ -53,7 +53,7 @@ python-openid==2.2.5
 pytz==2017.2
 pyzipcode==1.0
 redis==2.10.6
-requests==-2.23.0
+requests==2.23.0
 speaklater==1.3
 SQLAlchemy==0.8.0
 sqlalchemy-migrate==0.7.2