Dodat import/export Kontrola

backend/core/admin.py

@@ -1,5 +1,5 @@
 from django.contrib import admin
-from .models import Document, DocumentSegment, Organization, Risk
+from .models import Document, DocumentSegment, Organization, Risk, Control
 
 class DocumentSegmentInline(admin.StackedInline):
     model = DocumentSegment
@@ -21,7 +21,13 @@ class RiskAdmin(admin.ModelAdmin):
     ordering = ['risk_id']
     list_display = ['risk_id','risk_name','category']
 
+class ControlAdmin(admin.ModelAdmin):
+    list_display = ('risk', 'safeguard', 'weight')
+    search_fields = ('risk__risk_name', 'safeguard')
+    list_filter = ('risk', 'weight')
+
 
 admin.site.register(Document, DocumentAdmin)
 admin.site.register(Organization, OrganizationAdmin)
 admin.site.register(Risk ,RiskAdmin)
+admin.site.register(Control, ControlAdmin)

backend/core/management/commands/export_controls.py

@@ -0,0 +1,33 @@
+import csv
+from django.core.management.base import BaseCommand
+from backend.core.models import Control
+
+class Command(BaseCommand):
+    help = "Export controls to CSV file"
+
+    def add_arguments(self, parser):
+        parser.add_argument("csv_file", type=str, help="CSV file to export controls to")
+
+    def handle(self, *args, **options):
+        csv_file_path = options["csv_file"]
+
+        with open(csv_file_path, mode="w", newline="", encoding="utf-8") as csv_file:
+            fieldnames = [
+                "Risk #",
+                "Risk Description",
+                "CIS v8.1 Safeguards (Sub-Controls)",
+                "Weight (0-10)",
+            ]
+            writer = csv.DictWriter(csv_file, fieldnames=fieldnames)
+
+            writer.writeheader()
+
+            for control in Control.objects.select_related("risk").all():
+                writer.writerow({
+                    "Risk #": control.risk.risk_id,
+                    "Risk Description": control.risk.risk_name,
+                    "CIS v8.1 Safeguards (Sub-Controls)": control.safeguard,
+                    "Weight (0-10)": control.weight,
+                })
+
+        self.stdout.write(self.style.SUCCESS(f"Controls exported successfully to {csv_file_path}"))

backend/core/management/commands/import_controls.py

@@ -0,0 +1,38 @@
+import csv
+from django.core.management.base import BaseCommand
+from backend.core.models import Risk, Control
+
+class Command(BaseCommand):
+    help = "Import controls from CSV file"
+
+    def add_arguments(self, parser):
+        parser.add_argument("csv_file", type=str, help="CSV file to import controls from")
+
+    def handle(self, *args, **options):
+        csv_file_path = options["csv_file"]
+
+        with open(csv_file_path, mode="r", encoding="utf-8") as csv_file:
+            reader = csv.DictReader(csv_file)
+
+            for row in reader:
+                risk_id = int(row["Risk #"].strip())  
+                risk_desc = row["Risk Description"].strip() 
+                safeguard = row["CIS v8.1 Safeguards (Sub-Controls)"].strip()
+                
+                weight = int(row["Weight (0-10)"].strip()) if row["Weight (0-10)"].strip().isdigit() else 0 
+
+                risk, created = Risk.objects.get_or_create(
+                    risk_id=risk_id,
+                    defaults={"risk_name": risk_desc},
+                )
+
+                Control.objects.update_or_create(
+                    risk=risk,
+                    safeguard=safeguard,
+                    defaults={
+                        "description": risk_desc,
+                        "weight": weight,
+                    },
+                )
+
+        self.stdout.write(self.style.SUCCESS("Controls imported successfully!"))

backend/core/management/commands/tests/test_export_controls.py

@@ -0,0 +1,63 @@
+import os
+import csv
+from django.core.management import call_command
+from django.test import TestCase
+from backend.core.models import Risk, Control
+
+class ExportControlsCommandTest(TestCase):
+    def setUp(self):
+        self.risk1 = Risk.objects.create(risk_id=1, risk_name="Test Risk 1")
+        self.risk2 = Risk.objects.create(risk_id=2, risk_name="Test Risk 2")
+        
+        self.control1 = Control.objects.create(
+            risk=self.risk1,
+            safeguard="Test Safeguard 1",
+            description="Description 1",
+            weight=5
+        )
+        self.control2 = Control.objects.create(
+            risk=self.risk2,
+            safeguard="Test Safeguard 2",
+            description="Description 2",
+            weight=10
+        )
+
+        self.csv_file_path = 'test_export_controls.csv'
+
+    def tearDown(self):
+        if os.path.exists(self.csv_file_path):
+            os.remove(self.csv_file_path)
+
+    def test_export_controls_success(self):
+        call_command('export_controls', self.csv_file_path)
+
+        self.assertTrue(os.path.exists(self.csv_file_path))
+
+        with open(self.csv_file_path, mode='r', encoding='utf-8') as csv_file:
+            reader = csv.DictReader(csv_file)
+            rows = list(reader)
+
+            self.assertEqual(len(rows), 2) 
+
+            self.assertEqual(rows[0]["Risk #"], str(self.risk1.risk_id))
+            self.assertEqual(rows[0]["Risk Description"], self.risk1.risk_name)
+            self.assertEqual(rows[0]["CIS v8.1 Safeguards (Sub-Controls)"], self.control1.safeguard)
+            self.assertEqual(rows[0]["Weight (0-10)"], str(self.control1.weight))
+
+            self.assertEqual(rows[1]["Risk #"], str(self.risk2.risk_id))
+            self.assertEqual(rows[1]["Risk Description"], self.risk2.risk_name)
+            self.assertEqual(rows[1]["CIS v8.1 Safeguards (Sub-Controls)"], self.control2.safeguard)
+            self.assertEqual(rows[1]["Weight (0-10)"], str(self.control2.weight))
+
+    def test_export_controls_empty(self):
+        Control.objects.all().delete()
+
+        call_command('export_controls', self.csv_file_path)
+
+        self.assertTrue(os.path.exists(self.csv_file_path))
+
+        with open(self.csv_file_path, mode='r', encoding='utf-8') as csv_file:
+            reader = csv.DictReader(csv_file)
+            rows = list(reader)
+
+            self.assertEqual(len(rows), 0) 

backend/core/management/commands/tests/test_import_controls.py

@@ -0,0 +1,56 @@
+import os
+import csv
+import tempfile
+from django.core.management import call_command
+from django.test import TestCase
+from backend.core.models import Risk, Control
+
+class ImportControlsCommandTest(TestCase):
+    
+    def setUp(self):
+        """Create a temporary CSV file with test control data"""
+        self.temp_csv = tempfile.NamedTemporaryFile(mode="w", delete=False, newline='', encoding='utf-8')
+        fieldnames = [
+            "Risk #", "Risk Description", "CIS v8.1 Safeguards (Sub-Controls)", "Weight (0-10)"
+        ]
+        
+        writer = csv.DictWriter(self.temp_csv, fieldnames=fieldnames)
+        writer.writeheader()
+        writer.writerow({
+            "Risk #": "1",
+            "Risk Description": "Test Risk 1",
+            "CIS v8.1 Safeguards (Sub-Controls)": "Test Safeguard 1",
+            "Weight (0-10)": "5"
+        })
+        writer.writerow({
+            "Risk #": "2",
+            "Risk Description": "Test Risk 2",
+            "CIS v8.1 Safeguards (Sub-Controls)": "Test Safeguard 2",
+            "Weight (0-10)": "10"
+        })
+        self.temp_csv.close()
+
+    def test_import_controls_command(self):
+        """Test importing controls from a CSV file"""
+        self.assertEqual(Control.objects.count(), 0)
+        
+        self.assertEqual(Risk.objects.count(), 0)
+
+        Risk.objects.create(risk_id=1, risk_name="Test Risk 1")
+        Risk.objects.create(risk_id=2, risk_name="Test Risk 2")
+
+        call_command('import_controls', self.temp_csv.name)
+        
+        self.assertEqual(Control.objects.count(), 2)
+
+        control1 = Control.objects.get(risk__risk_id=1)
+        self.assertEqual(control1.safeguard, "Test Safeguard 1")
+        self.assertEqual(control1.weight, 5)
+
+        control2 = Control.objects.get(risk__risk_id=2)
+        self.assertEqual(control2.safeguard, "Test Safeguard 2")
+        self.assertEqual(control2.weight, 10)
+
+    def tearDown(self):
+        """Remove temporary CSV file after test"""
+        os.remove(self.temp_csv.name)

backend/core/migrations/0005_control.py

@@ -0,0 +1,24 @@
+# Generated by Django 5.1.3 on 2025-02-12 13:42
+
+import django.db.models.deletion
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('core', '0004_alter_document_id'),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='Control',
+            fields=[
+                ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('description', models.TextField()),
+                ('safeguard', models.CharField(max_length=255)),
+                ('weight', models.IntegerField()),
+                ('risk', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='controls', to='core.risk')),
+            ],
+        ),
+    ]

backend/core/models.py

@@ -141,3 +141,11 @@ class Risk(models.Model):
         return f"{self.risk_id} - {self.risk_name}"
 
 
+class Control(models.Model):
+    risk = models.ForeignKey(Risk, on_delete=models.CASCADE, related_name="controls")
+    description = models.TextField()
+    safeguard = models.CharField(max_length=255)
+    weight = models.IntegerField()
+    
+    def __str__(self):
+        return f"{self.safeguard} ({self.weight})"