diff --git a/backend/accounts/tasks.py b/backend/accounts/tasks.py
index 7ed1aa0..56cf39a 100644
--- a/backend/accounts/tasks.py
+++ b/backend/accounts/tasks.py
@@ -1,63 +1,52 @@
from celery import shared_task
from backend.core.models import Organization, Document, Risk, Control, DocumentRiskControl
from backend.core.utils import get_top_risk, get_controls_for_risk
-from django.shortcuts import get_object_or_404
-from django.template import Template, Context
-import yaml
-from backend.core.models import DocumentTemplate
-from .utils import send_payment_email
+from django.shortcuts import get_object_or_404, render
+
@shared_task
def create_document_for_organization(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)
- template = DocumentTemplate.objects.first()
- template_content = template.content
+ document = Document.objects.create(organization=organization)
+ document.add_segment('h1', "Top 10 Risks Identified")
- risks_data = []
+ 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_data = []
- selected_controls = get_controls_for_risk(risk, organization)
+ controls_content += f"Risk: {risk.risk_id} - {risk.risk_name}\n"
+ selected_controls = get_controls_for_risk(risk ,organization=organization)
+
for control_id, weight in selected_controls:
- control = Control.objects.filter(id=control_id).first()
+ control = Control.objects.filter(id=control_id).first()
if control:
- document = Document.objects.create(organization=organization)
DocumentRiskControl.objects.create(
document=document,
risk=risk,
control=control,
- weight=weight,
+ weight=weight
)
-
- controls_data.append({
- 'name': control.name,
- 'weight': weight,
- })
-
- risks_data.append({
- 'name': risk.risk_name,
- 'description': risk.description,
- 'controls': controls_data,
- })
-
- context = {
- 'organization': organization,
- 'risks': risks_data,
- }
-
- django_template = Template(template_content)
- rendered_content = django_template.render(Context(context))
-
- document = Document.objects.create(organization=organization, content=rendered_content)
-
- document.add_segment('title', f"{organization.name} Risk Assessment Report", 1)
- document.add_segment('introduction', f"This document outlines the risks and controls for {organization.name}.", 2)
- document.add_segment('body', rendered_content, 3)
-
- send_payment_email(confirmation_email)
-
- return document
+ controls_content += f" - Control: {control.name} (Impact Weight: {weight}/10)\n"
+
+ controls_content += "\n"
+
+ document.add_segment('body', controls_content)
\ No newline at end of file
diff --git a/backend/accounts/views.py b/backend/accounts/views.py
index 45500d9..3c75431 100644
--- a/backend/accounts/views.py
+++ b/backend/accounts/views.py
@@ -17,10 +17,10 @@ def confirm_email(request, uuid):
confirmation = get_object_or_404(EmailConfirmation, uuid=uuid)
if confirmation.is_expired():
- return render(request, 'accounts/confirmation_expired.html', {'email': confirmation.email})
+ return render(request, 'confirmation_expired.html', {'email': confirmation.email})
task = create_document_for_organization.delay(confirmation.email)
- return render(request, 'accounts/confirmation_success.html',{'email':confirmation.email})
+ return render(request, 'confirmation_success.html', {'email': confirmation.email})
def resend_confirmation(request,email):
if request.method == 'POST':
diff --git a/backend/core/templates/document_detail.html b/backend/core/templates/document_detail.html
deleted file mode 100644
index 0c94d53..0000000
--- a/backend/core/templates/document_detail.html
+++ /dev/null
@@ -1,19 +0,0 @@
-{% extends "base.html" %}
-
-{% block content %}
-
-
Document Preview
-
Organization: {{ organization.name }}
-
{{ created_at }}
-
-
Identified Risks and Controls
-
- {% for risk, controls in risks_with_controls.items %}
- -
-
{{ risk.risk_name }}:
- {{ controls }}
-
- {% endfor %}
-
-
-{% endblock %}
diff --git a/backend/core/urls.py b/backend/core/urls.py
index 146096f..47bc788 100644
--- a/backend/core/urls.py
+++ b/backend/core/urls.py
@@ -12,5 +12,4 @@ urlpatterns = [
path('document//', v.document, name='document'),
path('preview//', v.template_preview, name='template_preview'),
path("payment/", v.payment_page, name="payment_page"),
- path('documentview//', v.docprew, name='generate_document_view'),
]
diff --git a/backend/core/utils.py b/backend/core/utils.py
index 8e67dad..7e164ca 100644
--- a/backend/core/utils.py
+++ b/backend/core/utils.py
@@ -1,17 +1,19 @@
from openai import OpenAI
from django.conf import settings
-from .models import Risk, Control, Document, DocumentRiskControl
+from .models import Risk, Control
import time
-def extract_risk_factors(organization):
- excluded_fields={"name","email"}
+def extract_organization_details(organization):
+ excluded_fields = {"name", "email"}
risk_data = {}
for field in organization._meta.get_fields():
if field.name not in excluded_fields and hasattr(organization, field.name):
value = getattr(organization, field.name)
if value:
- risk_data[field.name] = value
+ help_text = getattr(field, 'help_text', '').strip()
+ key = help_text if help_text else field.name
+ risk_data[key] = value
return risk_data
def get_top_risk(organization):
@@ -26,16 +28,21 @@ def get_top_risk(organization):
Category: {risk.category}
Name: {risk.risk_name}
Primary Impact: {risk.primary_impact}
+ Secondary Impact: {risk.secondary_impact}
+ Tertiary Impact: {risk.tretiary_impact}
+ Detection Difficulty: {risk.detection_difficulty}
+ Recovery Complexity: {risk.recovery_complexity}
+ Business Impact Severity: {risk.businnes_impact_severity}
""")
- risk_factors = extract_risk_factors(organization)
+ organization_details = extract_organization_details(organization)
prompt = f"""
You are an AI risk assessor. Based on the following company details and list of known risks,
identify the 10 most critical risks for this company. Respond only with risk IDs.
Company Details:
- {risk_factors}
+ {organization_details}
List of Risks:
{risk_list}
@@ -56,16 +63,20 @@ def get_top_risk(organization):
def get_controls_for_risk(risk, organization):
client = OpenAI(api_key=settings.OPENAI_API_KEY)
all_controls = Control.objects.all()
- control_list = []
-
- risk_factors = extract_risk_factors(organization)
+ organization_details = extract_organization_details(organization)
+ control_list = [f"Control ID: {control.id}, Control Name: {control.name}" for control in all_controls]
valid_control_ids = {control.id for control in all_controls}
+ control_map = {control.id: control.name for control in all_controls}
- for control in all_controls:
- control_list.append(f"Control ID: {control.id}, Control Name: {control.name}")
+ def fetch_controls(prompt):
+ response = client.chat.completions.create(
+ model="gpt-4o-mini",
+ messages=[{"role": "system", "content": prompt}]
+ )
+ return response.choices[0].message.content.strip()
prompt = f"""
- You are an expert in cybersecurity risk management. Given the risk "{risk.risk_name}" and its associated factors "{risk_factors}",
+ You are an expert in cybersecurity risk management. Given the risk "{risk.risk_name}" and its associated organization details "{organization_details}",
your task is to select **exactly 10 unique controls** from the provided list that best mitigate this risk. Each control should be assigned a weight between **1 and 10** based on its effectiveness in reducing the risk.
### Rules:
1. **Each control ID must be unique** (no duplicates).
@@ -77,68 +88,97 @@ def get_controls_for_risk(risk, organization):
{control_list}
### Expected Response Format (STRICTLY FOLLOW THIS FORMAT):
- ```
:
:
-
- ```
-
### Example Correct Response (NO DUPLICATES):
- ```
12 : 8
45 : 7
-
- ```
-
⚠️ **If you provide duplicate control IDs, your response will be rejected. Ensure all control IDs are unique.**
⚠️ **Follow the response format exactly. Any deviation will be considered invalid.**
"""
- for attempt in range(5):
- response = client.chat.completions.create(
- model="gpt-4o-mini",
- messages=[{"role": "system", "content": prompt}]
- )
- result = response.choices[0].message.content.strip()
- print(f"AI Response (Attempt {attempt+1}):\n{result}")
+ selected_controls = []
+ control_ids_seen = set()
- selected_controls = []
- valid = True
- control_ids_seen = set()
+ result = fetch_controls(prompt)
+
+ for line in result.split("\n"):
+ line = line.strip()
+ parts = line.split(":")
+ 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()
+ print(f"Control:{control_id_str} Weight:{weight_str}")
+ print(f"ControlType: {type(control_id_str)} WeightType: {type(weight_str)}")
+ control_id_str = ''.join(filter(str.isdigit, control_id_str))
+ weight_str = ''.join(filter(str.isdigit, weight_str))
+
+ if control_id_str and weight_str:
+ try:
+ control_id = int(control_id_str)
+ weight = int(weight_str)
+
+ if control_id in valid_control_ids and 1 <= weight <= 10 and control_id not in control_ids_seen:
+ selected_controls.append((control_id, weight))
+ control_ids_seen.add(control_id)
+ except ValueError:
+ continue
+
+ if len(selected_controls) == 10:
+ return selected_controls
+
+ while len(selected_controls) < 10:
+ missing_count = 10 - len(selected_controls)
+ remaining_controls = valid_control_ids - control_ids_seen
+ remaining_controls_list = [f"Control ID: {cid}, Control Name: {control_map[cid]}" for cid in remaining_controls]
+
+ retry_prompt = f"""
+ You are an expert in cybersecurity risk management. Given the risk "{risk.risk_name}" and its associated organization details "{organization_details}",
+ your task is to select **exactly {missing_count} unique controls** from the provided list that best mitigate this risk. Each control should be assigned a weight between **1 and 10** based on its effectiveness in reducing the risk.
+ ### Rules:
+ 1. **Each control ID must be unique** (no duplicates).
+ 2. **Only return control IDs and weights** in the exact format below.
+ 3. **Weights must be between 1 and 10** (1 = low impact, 10 = high impact).
+ 4. **Do NOT add explanations, descriptions, or extra text.**
+ 5. **Ensure that control IDs are randomly distributed and diverse across different categories.**
+ ### Available Controls:
+ {remaining_controls_list}
+
+ ### Expected Response Format (STRICTLY FOLLOW THIS FORMAT):
+ :
+ :
+ ### Example Correct Response (NO DUPLICATES):
+ 12 : 8
+ 45 : 7
+ ⚠️ **If you provide duplicate control IDs, your response will be rejected. Ensure all control IDs are unique.**
+ ⚠️ **Follow the response format exactly. Any deviation will be considered invalid.**
+ """
+
+ result = fetch_controls(retry_prompt)
for line in result.split("\n"):
line = line.strip()
-
parts = line.split(":")
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()
-
+ print(f"Control:{control_id} Weight:{weight_str}")
+ print(f"ControlType: {type(control_id)} WeightType: {type(weight_str)}")
+
control_id_str = ''.join(filter(str.isdigit, control_id_str))
weight_str = ''.join(filter(str.isdigit, weight_str))
-
+
if control_id_str and weight_str:
- control_id = int(control_id_str)
- weight = int(weight_str)
-
- if control_id in valid_control_ids and 1 <= weight <= 10:
- if control_id in control_ids_seen:
- valid = False
- break
- selected_controls.append((control_id, weight))
- control_ids_seen.add(control_id)
- else:
- valid = False
- break
-
- if valid and len(selected_controls) == 10:
- return selected_controls
-
- print("Invalid response or duplicate control IDs found, retrying...\n")
- time.sleep(2)
-
- print("Failed to get a valid response after multiple attempts.")
- return []
-
-
+ try:
+ control_id = int(control_id_str)
+ weight = int(weight_str)
+
+ if control_id in valid_control_ids and 1 <= weight <= 10 and control_id not in control_ids_seen:
+ selected_controls.append((control_id, weight))
+ control_ids_seen.add(control_id)
+ except ValueError:
+ continue
+ if not remaining_controls:
+ break
+ return selected_controls if len(selected_controls) == 10 else []
diff --git a/backend/core/views.py b/backend/core/views.py
index 08f46b5..bd88f14 100644
--- a/backend/core/views.py
+++ b/backend/core/views.py
@@ -1,13 +1,13 @@
import logging
-import yaml
from django.shortcuts import render, redirect , get_object_or_404
from .forms import OrganizationForm
-from .models import Organization,Document,Risk, DocumentTemplate,DocumentRiskControl
-from backend.accounts.utils import send_confirmation_email,send_document_email
+from .models import Organization,Document,Risk, DocumentTemplate
+from backend.core.utils import get_top_risk
+from django.urls import reverse
+from backend.accounts.utils import send_confirmation_email, send_document_email
from django.contrib.admin.views.decorators import staff_member_required
-
# @login_required
# def index(request):
# return HttpResponse('Django
Página simples.
')
@@ -26,6 +26,7 @@ def signup(request):
if form.is_valid():
form.save()
send_confirmation_email(form.data['email'])
+
return render(request, 'thankyou.html', {
'email': form.data['email'],
})
@@ -69,27 +70,3 @@ def payment_page(request):
return render(request, "payment.html", {"email": email})
-
-def docprew(request, document_id):
- doc = get_object_or_404(Document, id=document_id)
- org = doc.organization
-
- document_risk_controls = DocumentRiskControl.objects.filter(document=doc)
-
- unique_risks = Risk.objects.filter(id__in=document_risk_controls.values('risk_id')).distinct()
-
- risks_with_controls = {}
-
- for risk in unique_risks:
- related_controls = DocumentRiskControl.objects.filter(risk=risk)
-
- risk_controls = [control.control.name for control in related_controls]
-
- risks_with_controls[risk] = ", ".join(risk_controls)
-
- return render(request, 'document_detail.html', {
- 'document': doc,
- 'organization': org,
- 'created_at': doc.created_at,
- 'risks_with_controls': risks_with_controls,
- })
\ No newline at end of file