Add hooks for tags and comments (#353)

CERT-Polska · Apr 21, 2021 · 8af3dea · 8af3dea
1 parent 557bf38
commit 8af3dea
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 2 deletions.
diff --git a/mwdb/core/plugins.py b/mwdb/core/plugins.py
@@ -8,7 +8,7 @@
 from mwdb.core.config import app_config
 from mwdb.core.log import getLogger
 from mwdb.core.util import is_subdir
-from mwdb.model import Config, File, Object, TextBlob
+from mwdb.model import Comment, Config, File, Object, Tag, TextBlob
 
 logger = getLogger()
 
@@ -78,6 +78,18 @@ def on_created_text_blob(self, blob: TextBlob):
     def on_reuploaded_text_blob(self, blob: TextBlob):
         pass
 
+    @hook_handler_method
+    def on_created_tag(self, object: Object, tag: Tag):
+        pass
+
+    @hook_handler_method
+    def on_reuploaded_tag(self, object: Object, tag: Tag):
+        pass
+
+    @hook_handler_method
+    def on_created_comment(self, object: Object, comment: Comment):
+        pass
+
 
 class PluginHookHandler(PluginHookBase):
     def __init__(self):

diff --git a/mwdb/resources/comment.py b/mwdb/resources/comment.py
@@ -3,6 +3,7 @@
 from werkzeug.exceptions import NotFound
 
 from mwdb.core.capabilities import Capabilities
+from mwdb.core.plugins import hooks
 from mwdb.model import Comment, db
 from mwdb.schema.comment import CommentItemResponseSchema, CommentRequestSchema
 
@@ -121,6 +122,7 @@ def post(self, type, identifier):
         logger.info("comment added", extra={"comment": comment.object_id})
 
         db.session.refresh(comment)
+        hooks.on_created_comment(db_object, comment)
         schema = CommentItemResponseSchema()
         return schema.dump(comment)
 

diff --git a/mwdb/resources/tag.py b/mwdb/resources/tag.py
@@ -4,6 +4,7 @@
 from werkzeug.exceptions import NotFound
 
 from mwdb.core.capabilities import Capabilities
+from mwdb.core.plugins import hooks
 from mwdb.model import ObjectPermission, Tag, db, object_tag_table
 from mwdb.schema.tag import (
     TagItemResponseSchema,
@@ -179,10 +180,15 @@ def put(self, type, identifier):
             raise NotFound("Object not found")
 
         tag_name = obj["tag"]
-        db_object.add_tag(tag_name)
+        is_new = db_object.add_tag(tag_name)
 
         logger.info("Tag added", extra={"tag": tag_name, "dhash": db_object.dhash})
         db.session.refresh(db_object)
+        if is_new:
+            hooks.on_created_tag(db_object, tag_name)
+        else:
+            hooks.on_reuploaded_tag(db_object, tag_name)
+
         schema = TagItemResponseSchema(many=True)
         return schema.dump(db_object.tags)