#1073 wip

orochi/api/api.py

@@ -2,6 +2,7 @@
 
 from orochi.api.routers.auth import router as auth_router
 from orochi.api.routers.bookmarks import router as bookmarks_router
+from orochi.api.routers.customrules import router as customrules_router
 from orochi.api.routers.dumps import router as dumps_router
 from orochi.api.routers.folders import router as folders_router
 from orochi.api.routers.plugins import router as plugins_router
@@ -18,3 +19,4 @@
 api.add_router("/utils/", utils_router, tags=["Utils"])
 api.add_router("/bookmarks/", bookmarks_router, tags=["Bookmarks"])
 api.add_router("/rules/", rules_router, tags=["Rules"])
+api.add_router("/customrules/", customrules_router, tags=["Custom Rules"])

orochi/api/models.py

@@ -1,3 +1,4 @@
+from enum import Enum
 from typing import Dict, List, Optional
 
 from django.contrib.auth import get_user_model
@@ -6,9 +7,15 @@
 from ninja.orm import create_schema
 
 from orochi.website.defaults import OSEnum
-from orochi.website.models import Bookmark, Dump, Folder, Plugin
+from orochi.website.models import Bookmark, CustomRule, Dump, Folder, Plugin
 from orochi.ya.models import Rule
 
+
+class RULE_ACTION(str, Enum):
+    PUBLISH = "Publish"
+    UNPUBLISH = "Unpublish"
+
+
 ###################################################
 # Auth
 ###################################################
@@ -236,6 +243,21 @@ class BookmarksInSchema(Schema):
     query: Optional[str] = None
 
 
+###################################################
+# CustomRules
+###################################################
+class CustomRulesOutSchema(ModelSchema):
+    class Meta:
+        model = CustomRule
+        fields = ["id", "name", "path", "public", "user"]
+
+
+class CustomRuleEditInSchema(ModelSchema):
+    class Meta:
+        model = CustomRule
+        fields = ["public"]
+
+
 ###################################################
 # Rules
 ###################################################
@@ -252,3 +274,12 @@ class Meta:
 
 class ListStr(Schema):
     rule_ids: List[int]
+
+
+class ListStrAction(Schema):
+    rule_ids: List[int]
+    action: RULE_ACTION
+
+
+class RuleEditInSchena(Schema):
+    text: str

orochi/api/routers/customrules.py

@@ -0,0 +1,187 @@
+import os
+import shutil
+from typing import List
+
+from django.conf import settings
+from django.db.models import Q
+from django.http import HttpResponse
+from ninja import Query, Router
+from ninja.security import django_auth
+
+from orochi.api.filters import RulesFilter
+from orochi.api.models import (
+    RULE_ACTION,
+    CustomRulesOutSchema,
+    ErrorsOut,
+    ListStr,
+    ListStrAction,
+    SuccessResponse,
+)
+from orochi.website.models import CustomRule
+
+router = Router()
+
+
+@router.get("/", response={200: List[CustomRulesOutSchema]}, auth=django_auth)
+def list_custom_rules(request, filters: Query[RulesFilter]):
+    return CustomRule.objects.all()
+
+
+@router.post(
+    "/{int:id}/default",
+    auth=django_auth,
+    url_name="default_customrule",
+    response={200: SuccessResponse, 400: ErrorsOut},
+)
+def default_rule(request, id: int):
+    """
+    Set a custom rule as the default.
+
+    Args:
+        request: The request object.
+        id (int): The ID of the custom rule to set as default.
+
+    Returns:
+        tuple: A tuple containing the status code and a dictionary with a message.
+
+    Raises:
+        Exception: If an error occurs during the process of setting the rule as default.
+    """
+    try:
+        old_default = CustomRule.objects.filter(user=request.user, default=True)
+        if old_default.count() == 1:
+            old = old_default.first()
+            old.default = False
+            old.save()
+
+        rule = CustomRule.objects.get(pk=id)
+        name = os.path.basename(rule.path)
+        if rule.user == request.user:
+            rule.default = True
+            rule.save()
+            return 200, {"message": f"Rule {name} set as default."}
+        # Make a copy
+        user_path = f"{settings.LOCAL_YARA_PATH}/{request.user.username}-Ruleset"
+        os.makedirs(user_path, exist_ok=True)
+        new_path = f"{user_path}/{rule.name}"
+        filename, extension = os.path.splitext(new_path)
+        counter = 1
+        while os.path.exists(new_path):
+            new_path = f"{filename}{counter}{extension}"
+            counter += 1
+
+        shutil.copy(rule.path, new_path)
+        CustomRule.objects.create(
+            user=request.user, name=rule.name, path=new_path, default=True
+        )
+        name = os.path.basename(new_path)
+
+        return 200, {
+            "message": f"Rule {name} copied in your ruleset and set as default."
+        }
+    except Exception as excp:
+        return 400, {"errors": str(excp)}
+
+
+@router.post(
+    "/publish",
+    auth=django_auth,
+    url_name="publish_customrule",
+    response={200: SuccessResponse, 400: ErrorsOut},
+)
+def publish_custom_rules(request, info: ListStrAction):
+    try:
+        rules = CustomRule.objects.filter(pk__in=info.rule_ids, user=request.user)
+        rules_count = rules.count()
+        for rule in rules:
+            rule.public = info.action == RULE_ACTION.PUBLISH
+            rule.save()
+        return 200, {"message": f"{rules_count} custom rules {info.action}ed."}
+
+    except Exception as excp:
+        return 400, {
+            "errors": (str(excp) if excp else "Generic error during publishing")
+        }
+
+
+@router.get("/{int:id}/download", auth=django_auth)
+def download(request, id: int):
+    """
+    Download a custom rule file by its primary key.
+
+    Args:
+        pk (int): The primary key of the custom rule to download.
+
+    Returns:
+        HttpResponse: The HTTP response object containing the downloaded custom rule file.
+
+    Raises:
+        Exception: If an error occurs during the process.
+    """
+    try:
+        rule = CustomRule.objects.filter(pk=id).filter(
+            Q(user=request.user) | Q(public=True)
+        )
+        if rule.count() == 1:
+            rule = rule.first()
+        else:
+            return 400, {"errors": "Generic error"}
+        if os.path.exists(rule.path):
+            with open(rule.path, "rb") as f:
+                rule_data = f.read()
+
+            response = HttpResponse(
+                rule_data,
+                content_type="application/text",
+            )
+            response["Content-Disposition"] = (
+                f"attachment; filename={os.path.basename(rule.path)}"
+            )
+            return response
+        else:
+            return 400, {"errors": "Custom Rule not found"}
+    except Exception as excp:
+        return 400, {"errors": str(excp)}
+
+
+@router.delete(
+    "/",
+    auth=django_auth,
+    url_name="delete_customrules",
+    response={200: SuccessResponse, 400: ErrorsOut},
+)
+def delete_custom_rules(request, info: ListStr):
+    """
+    Summary:
+    Delete custom rules based on the provided rule IDs.
+
+    Explanation:
+    This function deletes custom rules based on the specified rule IDs belonging to the authenticated user. It removes the rules from the database and returns a success message upon deletion.
+
+    Args:
+    - request: The request object.
+    - rule_ids: A list of integers representing the IDs of custom rules to be deleted.
+
+    Returns:
+    - Tuple containing status code and a message dictionary.
+
+    Raises:
+    - Any exception encountered during the process will result in a 400 status code with an error message.
+    """
+    try:
+        rules = CustomRule.objects.filter(pk__in=info.rule_ids, user=request.user)
+        rules_count = rules.count()
+        for rule in rules:
+            os.remove(rule.path)
+        rules.delete()
+        delete_message = f"{rules_count} custom rules deleted."
+        if rules_count != len(info.rule_ids):
+            delete_message += " Only custom rules in your ruleset have been deleted."
+        return 200, {"message": delete_message}
+
+    except Exception as excp:
+        return 400, {
+            "errors": (
+                str(excp) if excp else "Generic error during custom rules deletion"
+            )
+        }

orochi/api/routers/rules.py

@@ -1,7 +1,9 @@
 import os
+from pathlib import Path
 from typing import List
 
 import yara
+from django.conf import settings
 from django.http import HttpResponse
 from django.shortcuts import get_object_or_404
 from ninja import Query, Router
@@ -12,11 +14,12 @@
     ErrorsOut,
     ListStr,
     RuleBuildSchema,
+    RuleEditInSchena,
     RulesOutSchema,
     SuccessResponse,
 )
 from orochi.website.models import CustomRule
-from orochi.ya.models import Rule
+from orochi.ya.models import Rule, Ruleset
 
 router = Router()
 
@@ -26,8 +29,53 @@ def list_rules(request, filters: Query[RulesFilter]):
     return Rule.objects.all()
 
 
-@router.get("/{pk}/download", auth=django_auth)
-def download(request, pk: int):
+@router.patch(
+    "/{int:id}",
+    auth=django_auth,
+    url_name="edit_rule",
+    response={200: SuccessResponse, 400: ErrorsOut},
+)
+def edit_rule(request, id: int, data: RuleEditInSchena):
+    """
+    Edit or create a rule based on the provided primary key.
+
+    Args:
+        pk (int): The primary key of the rule to edit or create.
+
+    Returns:
+        tuple: A tuple containing the HTTP status code and a message indicating the success or error.
+    Raises:
+        Exception: If an error occurs during the process.
+    """
+    try:
+        rule = get_object_or_404(Rule, pk=id)
+        name = os.path.basename(rule.path)
+        if rule.ruleset.user == request.user:
+            with open(rule.path, "w") as f:
+                f.write(data.text)
+            return 200, {"message": f"Rule {name} updated."}
+        ruleset = get_object_or_404(Ruleset, user=request.user)
+        user_path = f"{settings.LOCAL_YARA_PATH}/{request.user.username}-Ruleset"
+        os.makedirs(user_path, exist_ok=True)
+        rule.pk = None
+        rule.ruleset = ruleset
+        new_path = f"{user_path}/{Path(rule.path).name}"
+        filename, extension = os.path.splitext(new_path)
+        counter = 1
+        while os.path.exists(new_path):
+            new_path = f"{filename}{counter}{extension}"
+            counter += 1
+        with open(new_path, "w") as f:
+            f.write(data.text)
+        rule.path = new_path
+        rule.save()
+        return 200, {"message": f"Rule {name} created in local ruleset."}
+    except Exception as excp:
+        return 400, {"errors": str(excp)}
+
+
+@router.get("/{int:id}/download", url_name="download_rule", auth=django_auth)
+def download(request, id: int):
     """
     Download a rule file by its primary key.
 
@@ -41,7 +89,7 @@ def download(request, pk: int):
         Exception: If an error occurs during the process.
     """
     try:
-        rule = Rule.objects.filter(pk=pk).filter(ruleset__enabled=True)
+        rule = Rule.objects.filter(pk=id).filter(ruleset__enabled=True)
         if rule.count() == 1:
             rule = rule.first()
         else:
@@ -90,12 +138,13 @@ def delete_rules(request, info: ListStr):
     """
     try:
         rules = Rule.objects.filter(pk__in=info.rule_ids, ruleset__user=request.user)
-        rules.delete()
         rules_count = rules.count()
-        if rules_count == 0:
-            return 200, {"message": f"{rules_count} rules deleted."}
-        else:
-            return 200, {"message": "Only rules in your ruleset can be deleted."}
+        rules.delete()
+        delete_message = f"{rules_count} rules deleted."
+        if rules_count != len(info.rule_ids):
+            delete_message += " Only rules in your ruleset have been deleted."
+        return 200, {"message": delete_message}
+
     except Exception as excp:
         return 400, {
             "errors": str(excp) if excp else "Generic error during rules deletion"

orochi/templates/users/user_bookmarks.html

@@ -177,7 +177,7 @@
             error: function (data) {
               $.toast({
                 title: 'Bookmark status!',
-                content: data.message,
+                content: data.errors,
                 type: 'error',
                 delay: 5000
               });
@@ -220,7 +220,7 @@
         error: function (data) {
           $.toast({
             title: 'Bookmark status!',
-            content: data.message,
+            content: data.errors,
             type: 'error',
             delay: 5000
           });