Reduce complexity in critical and complex methods

rekono/arguments/parser.py

@@ -50,7 +50,7 @@ def enumeration(enumeration: Enumeration, accumulated: dict = {}) -> dict:
         Keyword.HOST.name.lower(): enumeration.host.address,
         Keyword.PORT.name.lower(): enumeration.port,
         Keyword.PORTS.name.lower(): [enumeration.port],
-        Keyword.URL.name.lower(): get_url(None, enumeration),
+        Keyword.URL.name.lower(): get_url(enumeration.host.address, enumeration),
     }
     if accumulated and Keyword.PORTS.name.lower() in accumulated:
         output[Keyword.PORTS.name.lower()] = accumulated[Keyword.PORTS.name.lower()]

rekono/defectdojo/api/constants.py

@@ -0,0 +1 @@
+DD_FINDING_DATE_FORMAT = '%Y-%m-%d'

rekono/defectdojo/views.py

@@ -0,0 +1,74 @@
+from defectdojo import uploader
+from defectdojo.exceptions import (EngagementIdNotFoundException,
+                                   InvalidEngagementIdException,
+                                   ProductIdNotFoundException)
+from defectdojo.serializers import EngagementSerializer
+from drf_spectacular.utils import extend_schema
+from rest_framework import status
+from rest_framework.decorators import action
+from rest_framework.response import Response
+from rest_framework.viewsets import GenericViewSet
+
+
+class DDScansViewSet(GenericViewSet):
+
+    def get_executions(self):
+        return []
+
+    @extend_schema(request=EngagementSerializer, responses={200: None})
+    @action(
+        detail=True,
+        methods=['POST'],
+        url_path='defect-dojo-scans',
+        url_name='defect-dojo-scans'
+    )
+    def defect_dojo_scans(self, request, pk):
+        serializer = EngagementSerializer(data=request.data)
+        if serializer.is_valid():
+            try:
+                uploader.upload_executions(
+                    self.get_executions(),
+                    serializer.validated_data.get('engagement_id'),
+                    serializer.validated_data.get('engagement_name'),
+                    serializer.validated_data.get('engagement_description')
+                )
+                return Response(status=status.HTTP_200_OK)
+            except (
+                ProductIdNotFoundException,
+                EngagementIdNotFoundException,
+                InvalidEngagementIdException
+            ) as ex:
+                return Response(str(ex), status=status.HTTP_400_BAD_REQUEST)
+        return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST)
+
+
+class DDFindingsViewSet(GenericViewSet):
+
+    def get_findings(self):
+        return []
+
+    @extend_schema(request=EngagementSerializer, responses={200: None})
+    @action(
+        detail=True,
+        methods=['POST'],
+        url_path='defect-dojo-findings',
+        url_name='defect-dojo-findings'
+    )
+    def defect_dojo_findings(self, request, pk):
+        serializer = EngagementSerializer(data=request.data)
+        if serializer.is_valid():
+            try:
+                uploader.upload_findings(
+                    self.get_findings(),
+                    serializer.validated_data.get('engagement_id'),
+                    serializer.validated_data.get('engagement_name'),
+                    serializer.validated_data.get('engagement_description')
+                )
+                return Response(status=status.HTTP_200_OK)
+            except (
+                ProductIdNotFoundException,
+                EngagementIdNotFoundException,
+                InvalidEngagementIdException
+            ) as ex:
+                return Response(str(ex), status=status.HTTP_400_BAD_REQUEST)
+        return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST)

rekono/executions/migrations/0001_initial.py

@@ -1,15 +1,13 @@
-# Generated by Django 3.2.7 on 2021-11-06 16:18
+# Generated by Django 3.2.7 on 2021-11-06 20:15
 
 from django.db import migrations, models
-import django.db.models.deletion
 
 
 class Migration(migrations.Migration):
 
     initial = True
 
     dependencies = [
-        ('processes', '0002_initial'),
     ]
 
     operations = [
@@ -26,7 +24,6 @@ class Migration(migrations.Migration):
                 ('start', models.DateTimeField(blank=True, null=True)),
                 ('end', models.DateTimeField(blank=True, null=True)),
                 ('reported_to_defectdojo', models.BooleanField(default=False)),
-                ('step', models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.SET_NULL, to='processes.step')),
             ],
             options={
                 'ordering': ['-id'],

rekono/executions/migrations/0002_execution_step.py

@@ -0,0 +1,22 @@
+# Generated by Django 3.2.7 on 2021-11-06 20:15
+
+from django.db import migrations, models
+import django.db.models.deletion
+
+
+class Migration(migrations.Migration):
+
+    initial = True
+
+    dependencies = [
+        ('executions', '0001_initial'),
+        ('processes', '0001_initial'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='execution',
+            name='step',
+            field=models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.SET_NULL, to='processes.step'),
+        ),
+    ]

rekono/executions/migrations/0002_execution_task.py → rekono/executions/migrations/0003_execution_task.py

@@ -1,4 +1,4 @@
-# Generated by Django 3.2.7 on 2021-11-06 16:18
+# Generated by Django 3.2.7 on 2021-11-06 20:15
 
 from django.db import migrations, models
 import django.db.models.deletion
@@ -9,8 +9,8 @@ class Migration(migrations.Migration):
     initial = True
 
     dependencies = [
+        ('executions', '0002_execution_step'),
         ('tasks', '0001_initial'),
-        ('executions', '0001_initial'),
     ]
 
     operations = [

rekono/executions/queue/consumer.py

@@ -101,35 +101,32 @@ def get_new_jobs_from_findings(findings: dict, inputs: list) -> set:
         job_counter: []
     }
     for input_type in finding_relations.keys():
-        input_class = tool_utils.get_finding_class_by_type(input_type)
-        inputs = [i for i in inputs if i.type == input_type]
-        if not inputs or input_type not in findings:
+        if input_type not in findings:
             continue
-        for i in inputs:
+        input_class = tool_utils.get_finding_class_by_type(input_type)
+        for i in [i for i in inputs if i.type == input_type]:
             if finding_relations[input_type]:
                 relations_found = False
                 for finding in findings[input_type]:
                     for relation in finding_relations[input_type]:
-                        if hasattr(finding, relation.name.lower()):
-                            attribute = getattr(finding, relation.name.lower(), None)
-                            if attribute:
-                                relations_found = True
-                                for jc in jobs.copy():
-                                    if attribute in jobs[jc]:
-                                        if i.selection == InputSelection.ALL:
-                                            jobs[jc].append(finding)
+                        attribute = getattr(finding, relation.name.lower(), None)
+                        if attribute:
+                            relations_found = True
+                            for jc in jobs.copy():
+                                if attribute in jobs[jc]:
+                                    if i.selection == InputSelection.ALL:
+                                        jobs[jc].append(finding)
+                                    else:
+                                        related_items = [
+                                            f for f in jobs[jc] if not isinstance(f, input_class)
+                                        ]
+                                        if len(related_items) < len(jobs[jc]):
+                                            jobs[job_counter] = related_items.copy()
+                                            jobs[job_counter].append(finding)
+                                            job_counter += 1
                                         else:
-                                            related_items = [
-                                                f for f in jobs[jc]
-                                                if not isinstance(f, input_class)
-                                            ]
-                                            if len(related_items) < len(jobs[jc]):
-                                                jobs[job_counter] = related_items.copy()
-                                                jobs[job_counter].append(finding)
-                                                job_counter += 1
-                                            else:
-                                                jobs[jc].append(finding)
-                                break
+                                            jobs[jc].append(finding)
+                            break
                     if not relations_found:
                         for jc in jobs.copy():
                             jobs[jc].append(finding)

rekono/executions/views.py

@@ -1,28 +1,15 @@
-from defectdojo import uploader
-from defectdojo.exceptions import (EngagementIdNotFoundException,
-                                   InvalidEngagementIdException,
-                                   ProductIdNotFoundException)
-from defectdojo.serializers import EngagementSerializer
-from drf_spectacular.utils import extend_schema
+from defectdojo.views import DDFindingsViewSet, DDScansViewSet
 from executions.filters import ExecutionFilter
 from executions.models import Execution
 from executions.serializers import ExecutionSerializer
 from findings.models import (OSINT, Credential, Endpoint, Enumeration, Exploit,
                              Host, Technology, Vulnerability)
-from rest_framework import status
-from rest_framework.decorators import action
 from rest_framework.mixins import ListModelMixin, RetrieveModelMixin
-from rest_framework.response import Response
-from rest_framework.viewsets import GenericViewSet
 
 # Create your views here.
 
 
-class ExecutionViewSet(
-    GenericViewSet,
-    ListModelMixin,
-    RetrieveModelMixin
-):
+class ExecutionViewSet(ListModelMixin, RetrieveModelMixin, DDScansViewSet, DDFindingsViewSet):
     queryset = Execution.objects.all()
     serializer_class = ExecutionSerializer
     filterset_class = ExecutionFilter
@@ -31,66 +18,19 @@ def get_queryset(self):
         queryset = super().get_queryset()
         return queryset.filter(task__target__project__members=self.request.user)
 
-    @extend_schema(request=EngagementSerializer, responses={200: None})
-    @action(
-        detail=True,
-        methods=['POST'],
-        url_path='defect-dojo-scans',
-        url_name='defect-dojo-scans'
-    )
-    def defect_dojo_scans(self, request, pk):
-        execution = self.get_object()
-        serializer = EngagementSerializer(data=request.data)
-        if serializer.is_valid():
-            try:
-                uploader.upload_executions(
-                    [execution],
-                    serializer.validated_data.get('engagement_id'),
-                    serializer.validated_data.get('engagement_name'),
-                    serializer.validated_data.get('engagement_description')
-                )
-                return Response(status=status.HTTP_200_OK)
-            except (
-                ProductIdNotFoundException,
-                EngagementIdNotFoundException,
-                InvalidEngagementIdException
-            ) as ex:
-                return Response(str(ex), status=status.HTTP_400_BAD_REQUEST)
-        return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST)
+    def get_executions(self):
+        return [self.get_object()]
 
-    @extend_schema(request=EngagementSerializer, responses={200: None})
-    @action(
-        detail=True,
-        methods=['POST'],
-        url_path='defect-dojo-findings',
-        url_name='defect-dojo-findingss'
-    )
-    def defect_dojo_findings(self, request, pk):
+    def get_findings(self):
         execution = self.get_object()
-        serializer = EngagementSerializer(data=request.data)
-        if serializer.is_valid():
-            try:
-                findings = []
-                for find_model in [
-                    OSINT, Host, Enumeration, Technology,
-                    Endpoint, Vulnerability, Credential, Exploit
-                ]:
-                    findings.extend(find_model.objects.filter(
-                        execution=execution,
-                        is_active=True,
-                        is_manual=False
-                    ).all())
-                uploader.upload_findings(
-                    findings,
-                    serializer.validated_data.get('engagement_id'),
-                    serializer.validated_data.get('engagement_name'),
-                    serializer.validated_data.get('engagement_description')
-                )
-                return Response(status=status.HTTP_200_OK)
-            except (
-                ProductIdNotFoundException,
-                EngagementIdNotFoundException,
-                InvalidEngagementIdException
-            ) as ex:
-                return Response(str(ex), status=status.HTTP_400_BAD_REQUEST)
-        return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST)
+        findings = []
+        for find_model in [
+            OSINT, Host, Enumeration, Technology,
+            Endpoint, Vulnerability, Credential, Exploit
+        ]:
+            findings.extend(find_model.objects.filter(
+                execution=execution,
+                is_active=True,
+                is_manual=False
+            ).all())
+        return findings

rekono/findings/migrations/0001_initial.py

@@ -1,4 +1,4 @@
-# Generated by Django 3.2.7 on 2021-11-06 16:18
+# Generated by Django 3.2.7 on 2021-11-06 20:15
 
 from django.db import migrations, models
 import django.db.models.deletion

rekono/findings/migrations/0002_initial.py

@@ -1,4 +1,4 @@
-# Generated by Django 3.2.7 on 2021-11-06 16:18
+# Generated by Django 3.2.7 on 2021-11-06 20:15
 
 from django.db import migrations, models
 import django.db.models.deletion
@@ -9,9 +9,9 @@ class Migration(migrations.Migration):
     initial = True
 
     dependencies = [
-        ('findings', '0001_initial'),
         ('tasks', '0001_initial'),
-        ('executions', '0002_execution_task'),
+        ('findings', '0001_initial'),
+        ('executions', '0003_execution_task'),
     ]
 
     operations = [