Closes #1226 Adding Playbook Cache Feature (#1287)

* Adding /api/get_playbook_config API endpoint

* Adding codefactor suggestions to remove trailing whitespaces

* removing whitespaces

* Fixing schema bug

* Finishing up views, models and serializers with some active bugs left to take care of

* Removing unnecessary files

* Fixing linting issues

* Fixing API endpoints and making them work

* Fixing overall logic and laying down the endpoints more efficiently

* Fixing typo in serializer

* Fixing bugs wrt validation of playbooks when adding new ones

* Fixing linting errors

* Fixing flake8 errors

* Removing exposing responses i left out for debugging

* fixing unnecessary f strings

* Saving frontend progress

* Finishing up UI skeleton

* Fixing eslint issues in the frontend

* Running prettier on the frontend

* Adding code-doctor suggestions

* Fixing some bugs and making codedoctor changes

* Adding code factor suggestions

* Removing unnecessary printing

* running prettier on frontend

* Finishing up light test cases for the backend

* Finishing up test cases

* Finalising test cases

* Adding suggested changes

* removing pnpm-lock.yaml

* Fixing permissions

* Fixing typo in tests

* Moving test cases up for playbooks and fixing test cases

* Fixing test case errors

* Fixing typo

* Putting workflows back to the same place

* removing google file and fixed typo

* added docs about changes

Co-authored-by: Matteo Lodi <30625432+mlodic@users.noreply.github.com>

.github/workflows/pull_request_automation.yml

@@ -136,9 +136,9 @@ jobs:
         run: |
           docker/scripts/coverage_test.sh tests.analyzers_manager.test_file_scripts
 
-      - name: "Test: Playbooks Manager"
+      - name: "Test: Playbooks Manager (controllers, views)"
         run: |
-          docker/scripts/coverage_test.sh tests.playbooks_manager.test_controller
+          docker/scripts/coverage_test.sh tests.playbooks_manager.test_controller tests.playbooks_manager.test_views
 
       - name: "Coverage: generate xml and transfer from docker container to host"
         run: |

.gitignore

@@ -19,4 +19,5 @@ docker/.env.start
 .ipython
 .subversion
 .bash_history
-.cache
+.cache
+.google-cookie

api_app/migrations/0011_alter_organizationpluginstate_organization.py

@@ -0,0 +1,24 @@
+# Generated by Django 3.2.15 on 2022-11-12 13:51
+
+import django.db.models.deletion
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("certego_saas_organization", "0001_initial"),
+        ("api_app", "0010_custom_config_playbooks"),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name="organizationpluginstate",
+            name="organization",
+            field=models.ForeignKey(
+                on_delete=django.db.models.deletion.CASCADE,
+                related_name="+",
+                to="certego_saas_organization.organization",
+            ),
+        ),
+    ]

api_app/playbooks_manager/dataclasses.py

@@ -50,7 +50,7 @@ def get(cls, playbook_name: str) -> typing.Optional["PlaybookConfig"]:
         """
         Returns config dataclass by playbook_name if found, else None
         """
-        all_configs = cls.serializer_class.read_and_verify_config()
+        all_configs = cls.serializer_class.output_with_cached_playbooks()
         config_dict = all_configs.get(playbook_name, None)
         if config_dict is None:
             return None  # not found
@@ -60,7 +60,9 @@ def get(cls, playbook_name: str) -> typing.Optional["PlaybookConfig"]:
     def all(cls) -> typing.Dict[str, "PlaybookConfig"]:
         return {
             name: cls.from_dict(attrs)
-            for name, attrs in cls.serializer_class.read_and_verify_config().items()
+            for name, attrs in (
+                cls.serializer_class.output_with_cached_playbooks().items()
+            )
         }
 
     @classmethod

api_app/playbooks_manager/migrations/0001_initial.py

@@ -0,0 +1,42 @@
+# Generated by Django 3.2.15 on 2022-11-12 17:10
+
+import django.db.models.deletion
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    initial = True
+
+    dependencies = [
+        ("api_app", "0011_alter_organizationpluginstate_organization"),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name="CachedPlaybook",
+            fields=[
+                (
+                    "name",
+                    models.CharField(max_length=225, primary_key=True, serialize=False),
+                ),
+                (
+                    "description",
+                    models.CharField(blank=True, default="", max_length=225),
+                ),
+                ("analyzers", models.JSONField(default=dict)),
+                ("connectors", models.JSONField(default=dict)),
+                ("supports", models.JSONField(default=list)),
+                ("disabled", models.BooleanField(default=False)),
+                (
+                    "job",
+                    models.ForeignKey(
+                        on_delete=django.db.models.deletion.SET_NULL,
+                        null=True,
+                        to="api_app.job",
+                        related_name="job",
+                    ),
+                ),
+            ],
+        ),
+    ]

api_app/playbooks_manager/models.py

@@ -0,0 +1,23 @@
+# This file is a part of IntelOwl https://github.com/intelowlproject/IntelOwl
+# See the file 'LICENSE' for copying permission.
+
+from django.db import models
+
+from api_app.models import Job
+
+
+class CachedPlaybook(models.Model):
+    name = models.CharField(max_length=225, primary_key=True)
+    # Required fields
+    description = models.CharField(max_length=225, default="", blank=True)
+    analyzers = models.JSONField(default=dict)
+    connectors = models.JSONField(default=dict)
+
+    # Optional Fields
+    supports = models.JSONField(default=list)
+    disabled = models.BooleanField(default=False)
+
+    # job might not be necessary.
+    job = models.ForeignKey(
+        Job, on_delete=models.SET_NULL, related_name="job", null=True, blank=True
+    )

api_app/playbooks_manager/serializers.py

@@ -1,10 +1,22 @@
 # This file is a part of IntelOwl https://github.com/intelowlproject/IntelOwl
 # See the file 'LICENSE' for copying permission.
 
+import json
+import logging
+
+from django.core import serializers
 from rest_framework import serializers as rfs
 
 from api_app.analyzers_manager.constants import AllTypes
+from api_app.analyzers_manager.serializers import AnalyzerConfigSerializer
+from api_app.connectors_manager.serializers import ConnectorConfigSerializer
 from api_app.core.serializers import AbstractConfigSerializer
+from api_app.models import Job
+from certego_saas.apps.organization.permissions import IsObjectOwnerOrSameOrgPermission
+
+from .models import CachedPlaybook
+
+logger = logging.getLogger(__name__)
 
 
 class PlaybookConfigSerializer(AbstractConfigSerializer):
@@ -32,3 +44,155 @@ class PlaybookConfigSerializer(AbstractConfigSerializer):
         required=False,
         default=[],
     )
+
+    @classmethod
+    def _cached_playbooks(cls) -> dict:
+        """
+        Returns config file as `dict`.
+        """
+        config = super()._read_config()
+        cached_playbooks = CachedPlaybook.objects.all()
+        cached_serialized_playbooks = serializers.serialize(
+            "json", [obj for obj in cached_playbooks]
+        )
+        cached_playbooks_model_json = json.loads(cached_serialized_playbooks)
+        if len(cached_playbooks_model_json) == 0:
+            # this is for when no playbooks are cached
+            return config
+
+        cached_playbooks_final = {}
+        for playbook in cached_playbooks_model_json:
+            cached_playbooks_final[playbook["pk"]] = playbook["fields"]
+
+        return cached_playbooks_final
+
+    @classmethod
+    def output_with_cached_playbooks(cls) -> dict:
+        original_config_dict = cls.read_and_verify_config()
+        config_dict = cls._cached_playbooks()
+        serializer_errors = {}
+        for key, config in config_dict.items():
+            new_config = {"name": key, **config}
+            serializer = cls(data=new_config)  # lgtm [py/call-to-non-callable]
+            if serializer.is_valid():
+                config_dict[key] = serializer.data
+            else:
+                serializer_errors[key] = serializer.errors
+
+        if bool(serializer_errors):
+            logger.error(f"{cls.__name__} serializer failed: {serializer_errors}")
+            raise rfs.ValidationError(serializer_errors)
+
+        config_dict = config_dict | original_config_dict
+        return config_dict
+
+
+class CachedPlaybooksSerializer(rfs.ModelSerializer):
+    job_id = rfs.IntegerField()
+    name = rfs.CharField(max_length=225)
+    description = rfs.CharField(max_length=225)
+
+    class Meta:
+        model = CachedPlaybook
+        fields = (
+            "name",
+            "description",
+            "analyzers",
+            "connectors",
+            "supports",
+            "disabled",
+            "job_id",
+        )
+
+    def validate(self, attrs: dict) -> dict:
+        # The playbook in the playbook_config.json file is given more
+        # priority if the same named one is ever added back again.
+
+        attrs = super().validate(attrs)
+        playbook_name = attrs["name"].replace(" ", "_").upper()
+        job_id = attrs.get("job_id")
+
+        job = Job.objects.get(pk=job_id)
+        request = self.context.get("request", None)
+
+        has_perm = IsObjectOwnerOrSameOrgPermission().has_object_permission(
+            request, None, job
+        )
+
+        if not has_perm:
+            # bare in mind that for the time being,
+            # we don't check for which user is querying
+            # /api/get_playbook_configs and thus no filtering
+            # like this would be done on that end.
+            raise rfs.ValidationError(
+                "User doesn't have necessary permissions for this action."
+            )
+
+        analyzers_used = job.analyzers_to_execute
+        connectors_used = job.connectors_to_execute
+
+        analyzers = {analyzer: {} for analyzer in analyzers_used}
+        connectors = {connector: {} for connector in connectors_used}
+
+        supports = []
+        existing_playbooks = PlaybookConfigSerializer.output_with_cached_playbooks()
+
+        existing_playbook = existing_playbooks.get(playbook_name, {})
+
+        if existing_playbook != {}:
+            raise rfs.ValidationError("Another playbook exists with that name.")
+
+        analyzer_config = AnalyzerConfigSerializer.read_and_verify_config()
+        for analyzer_ in analyzers:
+            analyzer_checked = analyzer_config.get(analyzer_)
+            if analyzer_checked is None:
+                logger.info(f"Invalid analyzer {analyzer_}")
+            type_ = analyzer_checked.get("type")
+            if type_ == "file":
+                if type_ in supports:
+                    continue
+                supports.append(type_)
+            else:
+                observable_supported = analyzer_checked.get("observable_supported")
+                for observable_type in observable_supported:
+                    if observable_type not in supports:
+                        supports.append(observable_type)
+
+        connector_config = ConnectorConfigSerializer.read_and_verify_config()
+        for connector_ in connectors:
+            connector_checked = connector_config.get(connector_)
+            if connector_checked is None:
+                logger.info(f"Invalid connector {connector_}")
+
+        attrs["analyzers"] = analyzers
+        attrs["connectors"] = connectors
+        attrs["supports"] = supports
+        attrs["name"] = playbook_name
+
+        return attrs
+
+    def create(self, validated_data: dict) -> CachedPlaybook:
+        playbook_name = validated_data.get("name")
+        analyzers = validated_data.get("analyzers")
+        connectors = validated_data.get("connectors")
+        supports = validated_data.get("supports")
+        playbook_description = validated_data.get("description")
+        job_id = validated_data.get("job_id")
+        disabled = validated_data.get("disabled")
+
+        job = Job.objects.filter(pk=job_id).first()
+
+        if job is None:
+            raise rfs.ValidationError(f"Job of {job_id} doesn't exist.")
+
+        playbook = self.Meta.model.objects.create(
+            name=playbook_name,
+            analyzers={analyzer: {} for analyzer in analyzers},
+            connectors={connector: {} for connector in connectors},
+            supports=supports,
+            description=playbook_description,
+            job=job,
+            disabled=disabled if type(disabled) == bool else False,
+        )
+
+        return playbook

api_app/playbooks_manager/urls.py

@@ -3,10 +3,16 @@
 
 from django.urls import path
 
-from .views import PlaybookListAPI, analyze_multiple_files, analyze_multiple_observables
+from .views import (
+    PlaybookListAPI,
+    analyze_multiple_files,
+    analyze_multiple_observables,
+    cache_playbook_view,
+)
 
 urlpatterns = [
     path("get_playbook_configs", PlaybookListAPI.as_view()),
     path("playbook/analyze_multiple_files", analyze_multiple_files),
     path("playbook/analyze_multiple_observables", analyze_multiple_observables),
+    path("playbook/cache_playbook", cache_playbook_view),
 ]

api_app/playbooks_manager/views.py

@@ -11,7 +11,10 @@
 from rest_framework.decorators import api_view
 from rest_framework.response import Response
 
-from api_app.playbooks_manager.serializers import PlaybookConfigSerializer
+from api_app.playbooks_manager.serializers import (
+    CachedPlaybooksSerializer,
+    PlaybookConfigSerializer,
+)
 from api_app.serializers import (
     PlaybookAnalysisResponseSerializer,
     PlaybookFileAnalysisSerializer,
@@ -29,6 +32,36 @@
 ]
 
 
+def _cache_playbook(request, serializer_class: CachedPlaybooksSerializer):
+    """
+    Cache playbook after a scan
+    """
+
+    serializer = serializer_class(data=request.data, context={"request": request})
+    serializer.is_valid(raise_exception=True)
+
+    serializer.save()
+
+    return serializer.data
+
+
+@api_view(["POST"])
+def cache_playbook_view(
+    request,
+):
+    logger.info(f"received request from {request.user}." f"Data: {request.data}.")
+
+    response = _cache_playbook(
+        request=request,
+        serializer_class=CachedPlaybooksSerializer,
+    )
+
+    return Response(
+        response,
+        status=status.HTTP_200_OK,
+    )
+
+
 def _multi_analysis_request_playbooks(
     request,
     serializer_class: Union[
@@ -89,7 +122,7 @@ class PlaybookListAPI(APIView):
     )
     def get(self, request, *args, **kwargs):
         try:
-            pc = self.serializer_class.read_and_verify_config()
+            pc = self.serializer_class.output_with_cached_playbooks()
             return Response(pc, status=status.HTTP_200_OK)
         except Exception as e:
             logger.exception(