diff --git a/backend/core/admin.py b/backend/core/admin.py index 0647ada..ead9a41 100644 --- a/backend/core/admin.py +++ b/backend/core/admin.py @@ -36,9 +36,8 @@ class RiskAdmin(admin.ModelAdmin): 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') + list_display = ('id', 'name') + admin.site.register(Document, DocumentAdmin) diff --git a/backend/core/management/commands/export_controls.py b/backend/core/management/commands/export_controls.py index 4d3fb05..98dd3f3 100644 --- a/backend/core/management/commands/export_controls.py +++ b/backend/core/management/commands/export_controls.py @@ -3,7 +3,7 @@ from django.core.management.base import BaseCommand from backend.core.models import Control class Command(BaseCommand): - help = "Export controls to CSV file" + help = "Export controls (only safeguard) to CSV file" def add_arguments(self, parser): parser.add_argument("csv_file", type=str, help="CSV file to export controls to") @@ -12,22 +12,14 @@ class Command(BaseCommand): 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)", - ] + fieldnames = ["CIS v8.1 Safeguards (Sub-Controls)"] writer = csv.DictWriter(csv_file, fieldnames=fieldnames) writer.writeheader() - for control in Control.objects.select_related("risk").all(): + for control in Control.objects.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, + "CIS v8.1 Safeguards (Sub-Controls)": control.name, }) - self.stdout.write(self.style.SUCCESS(f"Controls exported successfully to {csv_file_path}")) \ No newline at end of file + self.stdout.write(self.style.SUCCESS(f"Controls exported successfully to {csv_file_path}")) diff --git a/backend/core/management/commands/import_controls.py b/backend/core/management/commands/import_controls.py index bcde2b2..e201705 100644 --- a/backend/core/management/commands/import_controls.py +++ b/backend/core/management/commands/import_controls.py @@ -1,9 +1,9 @@ import csv from django.core.management.base import BaseCommand -from backend.core.models import Risk, Control +from backend.core.models import Control class Command(BaseCommand): - help = "Import controls from CSV file" + help = "Import controls from CSV file (only safeguard) and store as name" def add_arguments(self, parser): parser.add_argument("csv_file", type=str, help="CSV file to import controls from") @@ -15,24 +15,11 @@ class Command(BaseCommand): 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, - }, + name=safeguard, + defaults={"name": safeguard}, ) - self.stdout.write(self.style.SUCCESS("Controls imported successfully!")) \ No newline at end of file + self.stdout.write(self.style.SUCCESS("Safeguards imported successfully!")) diff --git a/backend/core/management/commands/tests/test_export_controls.py b/backend/core/management/commands/tests/test_export_controls.py index 3accbd4..685d9c0 100644 --- a/backend/core/management/commands/tests/test_export_controls.py +++ b/backend/core/management/commands/tests/test_export_controls.py @@ -2,25 +2,12 @@ import os import csv from django.core.management import call_command from django.test import TestCase -from backend.core.models import Risk, Control +from backend.core.models import 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 - ) + Control.objects.create(name="Test Safeguard 1") + Control.objects.create(name="Test Safeguard 2") self.csv_file_path = 'test_export_controls.csv' @@ -28,7 +15,7 @@ class ExportControlsCommandTest(TestCase): if os.path.exists(self.csv_file_path): os.remove(self.csv_file_path) - def test_export_controls_success(self): + def test_export_controls(self): call_command('export_controls', self.csv_file_path) self.assertTrue(os.path.exists(self.csv_file_path)) @@ -37,27 +24,7 @@ class ExportControlsCommandTest(TestCase): reader = csv.DictReader(csv_file) rows = list(reader) - self.assertEqual(len(rows), 2) + 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 + self.assertEqual(rows[0]["CIS v8.1 Safeguards (Sub-Controls)"], "Test Safeguard 1") + self.assertEqual(rows[1]["CIS v8.1 Safeguards (Sub-Controls)"], "Test Safeguard 2") \ 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 index 37d355b..3b51fdb 100644 --- a/backend/core/management/commands/tests/test_import_controls.py +++ b/backend/core/management/commands/tests/test_import_controls.py @@ -1,56 +1,40 @@ 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 +from backend.core.models import 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) + self.csv_file_path = 'test_import_controls.csv' + with open(self.csv_file_path, mode='w', encoding='utf-8') as csv_file: + writer = csv.DictWriter(csv_file, fieldnames=["CIS v8.1 Safeguards (Sub-Controls)"]) + writer.writeheader() + writer.writerow({"CIS v8.1 Safeguards (Sub-Controls)": "Test Safeguard 1"}) + writer.writerow({"CIS v8.1 Safeguards (Sub-Controls)": "Test Safeguard 2"}) def tearDown(self): - """Remove temporary CSV file after test""" - os.remove(self.temp_csv.name) \ No newline at end of file + if os.path.exists(self.csv_file_path): + os.remove(self.csv_file_path) + + def test_import_controls(self): + self.assertEqual(Control.objects.count(), 0) + + call_command('import_controls', self.csv_file_path) + + self.assertEqual(Control.objects.count(), 2) + + safeguards = Control.objects.values_list('name', flat=True) + self.assertIn("Test Safeguard 1", safeguards) + self.assertIn("Test Safeguard 2", safeguards) + + def test_import_controls_update(self): + Control.objects.create(name="Test Safeguard 1") + + call_command('import_controls', self.csv_file_path) + + self.assertEqual(Control.objects.count(), 2) + + safeguards = Control.objects.values_list('name', flat=True) + self.assertIn("Test Safeguard 1", safeguards) + self.assertIn("Test Safeguard 2", safeguards) \ No newline at end of file diff --git a/backend/core/migrations/0007_rename_safeguard_control_name_and_more.py b/backend/core/migrations/0007_rename_safeguard_control_name_and_more.py new file mode 100644 index 0000000..ec5efd1 --- /dev/null +++ b/backend/core/migrations/0007_rename_safeguard_control_name_and_more.py @@ -0,0 +1,35 @@ +# Generated by Django 5.1.3 on 2025-02-13 17:09 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('core', '0006_documenttemplate'), + ] + + operations = [ + migrations.RenameField( + model_name='control', + old_name='safeguard', + new_name='name', + ), + migrations.RemoveField( + model_name='control', + name='description', + ), + migrations.RemoveField( + model_name='control', + name='risk', + ), + migrations.RemoveField( + model_name='control', + name='weight', + ), + migrations.AlterField( + model_name='control', + name='id', + field=models.AutoField(primary_key=True, serialize=False), + ), + ] diff --git a/backend/core/models.py b/backend/core/models.py index c1531a5..dd89a01 100644 --- a/backend/core/models.py +++ b/backend/core/models.py @@ -153,10 +153,8 @@ class Risk(models.Model): 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() - + id = models.AutoField(primary_key=True) + name = models.CharField(max_length=255) + def __str__(self): - return f"{self.safeguard} ({self.weight})" \ No newline at end of file + return f"{self.id} ({self.name})" \ No newline at end of file