diff --git a/backend/accounts/views.py b/backend/accounts/views.py index ecf9c75..0eccaf7 100644 --- a/backend/accounts/views.py +++ b/backend/accounts/views.py @@ -5,7 +5,8 @@ from .models import EmailConfirmation from django.shortcuts import get_object_or_404, render from django.http import HttpResponse from backend.accounts.utils import send_confirmation_email - +from backend.core.utils import get_controls_for_risk, get_top_risk +from backend.core.models import Organization,Risk,Document,Control,DocumentRiskControl class SignUpView(CreateView): form_class = SignupForm @@ -13,12 +14,71 @@ class SignUpView(CreateView): template_name = 'accounts/signup.html' -def confirm_email(request,uuid): +from django.urls import reverse_lazy +from django.views.generic import CreateView +from django.shortcuts import get_object_or_404, render +from django.http import HttpResponse +from backend.accounts.forms import SignupForm +from backend.accounts.utils import send_confirmation_email +from backend.core.utils import get_controls_for_risk, get_top_risk +from backend.core.models import Organization, Risk, Document, Control, DocumentRiskControl + +class SignUpView(CreateView): + form_class = SignupForm + success_url = reverse_lazy('login') + template_name = 'accounts/signup.html' + + +def confirm_email(request, uuid): confirmation = get_object_or_404(EmailConfirmation, uuid=uuid) if confirmation.is_expired(): - return render(request,'confirmation_expired.html', {'email': confirmation.email}) + return render(request, 'confirmation_expired.html', {'email': confirmation.email}) + organization = get_object_or_404(Organization, email=confirmation.email) + + 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 Risks Identified") + + risk_content = "\n\n".join([ + f"Risk: {risk.risk_id} - {risk.risk_name} \n" + f"Category: {risk.category}\n" + f"Primary Impact: {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}") + + controls_content = "Mitigation Controls:\n\n" + + for risk in top_risks: + controls_content += f"Risk: {risk.risk_id} - {risk.risk_name}\n" + + selected_controls = get_controls_for_risk(risk) + + for control_id, weight in selected_controls: + control = Control.objects.filter(id=control_id).first() + if control: + DocumentRiskControl.objects.create( + document=document, + risk=risk, + control=control, + weight=weight + ) + controls_content += f" - Control: {control.name} (Impact Weight: {weight}/10)\n" + + controls_content += "\n" + + document.add_segment('body', controls_content) + return HttpResponse("Email is confirmed") def resend_confirmation(request,email): diff --git a/backend/core/admin.py b/backend/core/admin.py index ead9a41..87beba0 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, 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) diff --git a/backend/core/migrations/0008_documentriskcontrol.py b/backend/core/migrations/0008_documentriskcontrol.py new file mode 100644 index 0000000..b993d7e --- /dev/null +++ b/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')}, + }, + ), + ] diff --git a/backend/core/models.py b/backend/core/models.py index dd89a01..138a2ee 100644 --- a/backend/core/models.py +++ b/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})" \ No newline at end of file + 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') \ No newline at end of file diff --git a/backend/core/utils.py b/backend/core/utils.py index a572c89..b261e10 100644 --- a/backend/core/utils.py +++ b/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: 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] diff --git a/backend/core/views.py b/backend/core/views.py index 7b4ecc7..1eed1e7 100644 --- a/backend/core/views.py +++ b/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)