#9 AI bira kontrole,i dinamicki se updejtuje document

backend/core/admin.py

@@ -1,5 +1,5 @@
 from django.contrib import admin
-from .models import Document, DocumentSegment, Organization, Risk, Control, DocumentTemplate
+from .models import Document, DocumentSegment, Organization, Risk, Control, DocumentTemplate, DocumentRiskControl
 from django.urls import reverse
 from django.utils.html import format_html
 
@@ -38,6 +38,8 @@ class RiskAdmin(admin.ModelAdmin):
 class ControlAdmin(admin.ModelAdmin):
     list_display = ('id', 'name')
 
+class DocumentRiskControlAdmin(admin.ModelAdmin):
+    list_display = ('document', 'risk', 'control', 'weight')
 
 
 admin.site.register(Document, DocumentAdmin)
@@ -45,3 +47,4 @@ admin.site.register(Organization, OrganizationAdmin)
 admin.site.register(Risk ,RiskAdmin)
 admin.site.register(Control, ControlAdmin)
 admin.site.register(DocumentTemplate, DocumentTemplateAdmin)
+admin.site.register(DocumentRiskControl, DocumentRiskControlAdmin)

backend/core/migrations/0008_documentriskcontrol.py

@@ -0,0 +1,27 @@
+# Generated by Django 5.1.3 on 2025-02-14 13:09
+
+import django.db.models.deletion
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('core', '0007_rename_safeguard_control_name_and_more'),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='DocumentRiskControl',
+            fields=[
+                ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('weight', models.IntegerField()),
+                ('control', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to='core.control')),
+                ('document', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to='core.document')),
+                ('risk', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to='core.risk')),
+            ],
+            options={
+                'unique_together': {('document', 'risk', 'control')},
+            },
+        ),
+    ]

backend/core/models.py

@@ -157,4 +157,13 @@ class Control(models.Model):
     name = models.CharField(max_length=255)
 
     def __str__(self):
-        return f"{self.id} ({self.name})"
+        return f"{self.id} ({self.name})"
+    
+class DocumentRiskControl(models.Model):
+    document = models.ForeignKey(Document, on_delete=models.CASCADE)
+    risk = models.ForeignKey(Risk, on_delete=models.CASCADE)
+    control = models.ForeignKey(Control, on_delete=models.CASCADE)
+    weight = models.IntegerField()
+
+    class Meta:
+        unique_together = ('document', 'risk', 'control')

backend/core/utils.py

@@ -1,6 +1,7 @@
 from openai import OpenAI
 from django.conf import settings
-from .models import Risk
+from .models import Risk, Control, Document, DocumentRiskControl
+from django.shortcuts import get_object_or_404
 
 def extract_risk_factors(organization):
     excluded_fields={"name","email"}
@@ -50,3 +51,52 @@ def get_top_risk(organization):
     risk_ids = response.choices[0].message.content.strip().split(",")
 
     return [int(risk_id) for risk_id in risk_ids if risk_id.isdigit()]
+
+def get_controls_for_risk(risk):
+    client = OpenAI(api_key=settings.OPENAI_API_KEY)
+    all_controls = Control.objects.all()
+    control_list = []
+
+    for control in all_controls:
+        control_list.append(f"Control ID: {control.id}, Control Name: {control.name}")
+
+    prompt = f"""
+    You are a cyber security expert. For the risk '{risk.risk_name}', select 10 relevant controls
+    from the following list and assign a weight (1-10) based on how much they reduce risks.
+    Available Controls (only respond with control IDs and weights):
+    {control_list}
+    Respond only with control IDs (numbers) and their corresponding weights (1-10).
+    Format:
+    ID: <control_id> Weight: <weight>
+    Example:
+    1: 9
+    2: 6
+    3: 4
+    """
+
+    response = client.chat.completions.create(
+        model="gpt-4",
+        messages=[{"role": "system", "content": prompt}]
+    )
+
+    result = response.choices[0].message.content.strip()
+    selected_controls = []
+
+    for line in result.split("\n"):
+        line = line.strip()
+
+        parts = line.split("Weight:")
+        if len(parts) == 2:
+            control_id_str = parts[0].replace("ID:", "").replace("id:", "").replace("Id:", "").strip() 
+            weight_str = parts[1].strip().replace("Weight:", "").replace("weight:","").strip()
+
+            control_id_str = ''.join(filter(str.isdigit, control_id_str))
+            weight_str = ''.join(filter(str.isdigit, weight_str))
+            control_id = int(control_id_str)  
+            weight = int(weight_str)  
+            print(f"ID: {control_id}, Weight: {weight}")
+
+            control = Control.objects.filter(id=control_id).first()
+            if control:
+                selected_controls.append((control_id, weight))
+    return selected_controls[:10]  

backend/core/views.py

@@ -24,34 +24,11 @@ def signup(request):
     if request.method == 'POST':
         form = OrganizationForm(request.POST)
         if form.is_valid():
-            organization = form.save()
-            top_risk_ids = get_top_risk(organization)
-            top_risks = Risk.objects.filter(risk_id__in = top_risk_ids)
-
-            organization.risks.set(top_risks)
-
-            document = Document.objects.create(organization=organization)
-            document.add_segment('h1', "Top 10 Risk Identified")
-
-            risk_content = "\n\n".join([
-                f"Risk: {risk.risk_id} : {risk.risk_name} \n"
-                f"Category: {risk.category}\n"
-                f"Primary Impaact: {risk.primary_impact} \n"
-                f"Secondary Impact: {risk.secondary_impact}\n"
-                f"Tertiary Impact: {risk.tretiary_impact} \n"
-                f"Detection Difficulty: {risk.detection_difficulty} \n"
-                f"Recovery Complexity: {risk.recovery_complexity} \n"
-                f"Business Impact Severity: {risk.businnes_impact_severity}\n"
-                for risk in top_risks
-            ])
-            
-            document.add_segment('body',f"Identified Risks: \n\n{risk_content}")
-
+            form.save()
             send_confirmation_email(form.data['email'])
 
             return render(request, 'thankyou.html', {
                 'email': form.data['email'],
-                'document_link': reverse('core:document', args=[str(document.id)])
                 })
         else:
             logging.error(form.errors)