From 5ed323140fe2827c5f4c19a8bfccbd49a7217042 Mon Sep 17 00:00:00 2001 From: Amir Date: Mon, 17 Feb 2025 20:36:59 +0100 Subject: [PATCH] #10 AI bira controle,dodate je dummy payment page --- backend/accounts/tasks.py | 73 +++++++----- .../accounts/confirmation_expired.html | 6 +- .../accounts/confirmation_success.html | 12 +- backend/accounts/utils.py | 2 - backend/accounts/views.py | 9 +- backend/core/templates/document_detail.html | 19 +++ backend/core/templates/payment.html | 17 +-- backend/core/urls.py | 2 +- backend/core/utils.py | 110 ++++++++++++------ backend/core/views.py | 32 ++++- start_services.sh | 12 ++ 11 files changed, 204 insertions(+), 90 deletions(-) create mode 100644 backend/core/templates/document_detail.html create mode 100755 start_services.sh diff --git a/backend/accounts/tasks.py b/backend/accounts/tasks.py index a75c1cb..7ed1aa0 100644 --- a/backend/accounts/tasks.py +++ b/backend/accounts/tasks.py @@ -1,52 +1,63 @@ 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, render - +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 @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) - document = Document.objects.create(organization=organization) - document.add_segment('h1', "Top 10 Risks Identified") + template = DocumentTemplate.objects.first() + template_content = template.content - 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" - + risks_data = [] for risk in top_risks: - controls_content += f"Risk: {risk.risk_id} - {risk.risk_name}\n" + controls_data = [] + selected_controls = get_controls_for_risk(risk, organization) - selected_controls = get_controls_for_risk(risk) - 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_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 + + 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 diff --git a/backend/accounts/templates/accounts/confirmation_expired.html b/backend/accounts/templates/accounts/confirmation_expired.html index 7624b9a..8ca56b0 100644 --- a/backend/accounts/templates/accounts/confirmation_expired.html +++ b/backend/accounts/templates/accounts/confirmation_expired.html @@ -1,8 +1,10 @@ {% extends "base_login.html" %} - +{%block content%}

Link has expired!

{% csrf_token %} -
\ No newline at end of file + + +{% endblock %} \ No newline at end of file diff --git a/backend/accounts/templates/accounts/confirmation_success.html b/backend/accounts/templates/accounts/confirmation_success.html index 32a64d0..a26424f 100644 --- a/backend/accounts/templates/accounts/confirmation_success.html +++ b/backend/accounts/templates/accounts/confirmation_success.html @@ -1,5 +1,11 @@ -{% extends "base_login.html" %} +{% extends "base.html" %} +{%block content%} +
+
+

Email Confirmed!

+

Your email {{ email }} has been successfully verified.

+
+
-

Email Confirmed!

-

Your email {{ email }} has been successfully verified.

+{%endblock%} diff --git a/backend/accounts/utils.py b/backend/accounts/utils.py index a6bbefd..c9a6ba8 100644 --- a/backend/accounts/utils.py +++ b/backend/accounts/utils.py @@ -6,8 +6,6 @@ from django.conf import settings from django.utils.timezone import now from backend.core.models import Document, Organization - - def send_confirmation_email(email): confirmation, created = EmailConfirmation.objects.get_or_create(email=email) diff --git a/backend/accounts/views.py b/backend/accounts/views.py index 1757f37..45500d9 100644 --- a/backend/accounts/views.py +++ b/backend/accounts/views.py @@ -4,7 +4,7 @@ from backend.accounts.forms import SignupForm 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, send_payment_email +from backend.accounts.utils import send_confirmation_email from .tasks import create_document_for_organization class SignUpView(CreateView): @@ -17,13 +17,10 @@ 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, 'accounts/confirmation_expired.html', {'email': confirmation.email}) task = create_document_for_organization.delay(confirmation.email) - send_payment_email(confirmation.email) - - - return HttpResponse("Email is confirmed") + return render(request, 'accounts/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 new file mode 100644 index 0000000..0c94d53 --- /dev/null +++ b/backend/core/templates/document_detail.html @@ -0,0 +1,19 @@ +{% extends "base.html" %} + +{% block content %} +
+

Document Preview

+

Organization: {{ organization.name }}

+

{{ created_at }}

+ +

Identified Risks and Controls

+ +
+{% endblock %} diff --git a/backend/core/templates/payment.html b/backend/core/templates/payment.html index 22bd6c6..ab52439 100644 --- a/backend/core/templates/payment.html +++ b/backend/core/templates/payment.html @@ -1,11 +1,14 @@ {% extends 'base.html' %} {% block content %} -

Payment

-

Click the button below to pay and access your document.

- -
- {% csrf_token %} - -
+
+
+

Payment

+

Click the button below to pay and access your document.

+
+ {% csrf_token %} + +
+
+
{% endblock %} \ No newline at end of file diff --git a/backend/core/urls.py b/backend/core/urls.py index 9b5d2ea..146096f 100644 --- a/backend/core/urls.py +++ b/backend/core/urls.py @@ -12,5 +12,5 @@ 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 b261e10..8e67dad 100644 --- a/backend/core/utils.py +++ b/backend/core/utils.py @@ -1,7 +1,7 @@ from openai import OpenAI from django.conf import settings from .models import Risk, Control, Document, DocumentRiskControl -from django.shortcuts import get_object_or_404 +import time def extract_risk_factors(organization): excluded_fields={"name","email"} @@ -44,59 +44,101 @@ def get_top_risk(organization): """ response = client.chat.completions.create( - model="gpt-4", + model="gpt-4o-mini", messages=[{"role": "system", "content": prompt}] ) risk_ids = response.choices[0].message.content.strip().split(",") + print(f"Risks: {risk_ids}") return [int(risk_id) for risk_id in risk_ids if risk_id.isdigit()] -def get_controls_for_risk(risk): +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) + valid_control_ids = {control.id for control in all_controls} + 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): + You are an expert in cybersecurity risk management. Given the risk "{risk.risk_name}" and its associated factors "{risk_factors}", + 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). + 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: {control_list} - Respond only with control IDs (numbers) and their corresponding weights (1-10). - Format: - ID: Weight: - Example: - 1: 9 - 2: 6 - 3: 4 + + ### 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.** """ - response = client.chat.completions.create( - model="gpt-4", - messages=[{"role": "system", "content": prompt}] - ) + 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}") - result = response.choices[0].message.content.strip() - selected_controls = [] + selected_controls = [] + valid = True + control_ids_seen = set() - for line in result.split("\n"): - line = line.strip() + 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() + + 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 [] - 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 19528da..08f46b5 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 -from backend.core.utils import get_top_risk -from django.urls import reverse +from .models import Organization,Document,Risk, DocumentTemplate,DocumentRiskControl 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,7 +26,6 @@ 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,3 +68,28 @@ def payment_page(request): return redirect(document_link) 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 diff --git a/start_services.sh b/start_services.sh new file mode 100755 index 0000000..84d747e --- /dev/null +++ b/start_services.sh @@ -0,0 +1,12 @@ +#!/bin/bash + +echo "Starting Django server..." +python3 manage.py runserver & + +echo "Starting Celery worker..." +celery -A backend worker --loglevel=info & + +echo "Starting Redis server on port 6380..." +redis-server --port 6380 & + +wait