From 39c3e250ed62d019b8a8581bb2321fce57cd18b7 Mon Sep 17 00:00:00 2001 From: Amir Date: Wed, 12 Feb 2025 16:15:06 +0100 Subject: [PATCH] Dodat import/export Kontrola --- backend/core/admin.py | 8 ++- .../management/commands/export_controls.py | 33 ++++++++++ .../management/commands/import_controls.py | 38 +++++++++++ .../commands/tests/test_export_controls.py | 63 +++++++++++++++++++ .../commands/tests/test_import_controls.py | 56 +++++++++++++++++ backend/core/migrations/0005_control.py | 24 +++++++ backend/core/models.py | 8 +++ 7 files changed, 229 insertions(+), 1 deletion(-) create mode 100644 backend/core/management/commands/export_controls.py create mode 100644 backend/core/management/commands/import_controls.py create mode 100644 backend/core/management/commands/tests/test_export_controls.py create mode 100644 backend/core/management/commands/tests/test_import_controls.py create mode 100644 backend/core/migrations/0005_control.py diff --git a/backend/core/admin.py b/backend/core/admin.py index 50dc76d..d3b2467 100644 --- a/backend/core/admin.py +++ b/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) diff --git a/backend/core/management/commands/export_controls.py b/backend/core/management/commands/export_controls.py new file mode 100644 index 0000000..4d3fb05 --- /dev/null +++ b/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}")) \ No newline at end of file diff --git a/backend/core/management/commands/import_controls.py b/backend/core/management/commands/import_controls.py new file mode 100644 index 0000000..bcde2b2 --- /dev/null +++ b/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!")) \ No newline at end of file diff --git a/backend/core/management/commands/tests/test_export_controls.py b/backend/core/management/commands/tests/test_export_controls.py new file mode 100644 index 0000000..3accbd4 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/backend/core/management/commands/tests/test_import_controls.py b/backend/core/management/commands/tests/test_import_controls.py new file mode 100644 index 0000000..37d355b --- /dev/null +++ b/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) \ No newline at end of file diff --git a/backend/core/migrations/0005_control.py b/backend/core/migrations/0005_control.py new file mode 100644 index 0000000..f5e1ae1 --- /dev/null +++ b/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')), + ], + ), + ] diff --git a/backend/core/models.py b/backend/core/models.py index 76514e4..bc25609 100644 --- a/backend/core/models.py +++ b/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})" \ No newline at end of file