#10 AI bira controle,dodate je dummy payment page

This commit is contained in:
2025-02-17 20:36:59 +01:00
parent a097590b62
commit 5ed323140f
11 changed files with 204 additions and 90 deletions

View File

@@ -1,52 +1,63 @@
from celery import shared_task from celery import shared_task
from backend.core.models import Organization, Document, Risk, Control, DocumentRiskControl from backend.core.models import Organization, Document, Risk, Control, DocumentRiskControl
from backend.core.utils import get_top_risk, get_controls_for_risk 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 @shared_task
def create_document_for_organization(confirmation_email): def create_document_for_organization(confirmation_email):
organization = get_object_or_404(Organization, email=confirmation_email) organization = get_object_or_404(Organization, email=confirmation_email)
top_risk_ids = get_top_risk(organization) top_risk_ids = get_top_risk(organization)
top_risks = Risk.objects.filter(risk_id__in=top_risk_ids) top_risks = Risk.objects.filter(risk_id__in=top_risk_ids)
organization.risks.set(top_risks)
document = Document.objects.create(organization=organization) template = DocumentTemplate.objects.first()
document.add_segment('h1', "Top 10 Risks Identified") template_content = template.content
risk_content = "\n\n".join([ risks_data = []
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: 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: 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: if control:
document = Document.objects.create(organization=organization)
DocumentRiskControl.objects.create( DocumentRiskControl.objects.create(
document=document, document=document,
risk=risk, risk=risk,
control=control, control=control,
weight=weight weight=weight,
) )
controls_content += f" - Control: {control.name} (Impact Weight: {weight}/10)\n"
controls_data.append({
controls_content += "\n" 'name': control.name,
'weight': weight,
document.add_segment('body', controls_content) })
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

View File

@@ -1,8 +1,10 @@
{% extends "base_login.html" %} {% extends "base_login.html" %}
{%block content%}
<h2>Link has expired! </h2> <h2>Link has expired! </h2>
<form method="POST"> <form method="POST">
{% csrf_token %} {% csrf_token %}
<button type="submit">Resend link</button> <button type="submit">Resend link</button>
</form>
</form>
{% endblock %}

View File

@@ -1,5 +1,11 @@
{% extends "base_login.html" %} {% extends "base.html" %}
{%block content%}
<div class="out-risk-management">
<div class="risk-management">
<h1>Email Confirmed!</h1>
<p>Your email {{ email }} has been successfully verified.</p>
</div>
</div>
<h1>Email Confirmed!</h1> {%endblock%}
<p>Your email {{ email }} has been successfully verified.</p>

View File

@@ -6,8 +6,6 @@ from django.conf import settings
from django.utils.timezone import now from django.utils.timezone import now
from backend.core.models import Document, Organization from backend.core.models import Document, Organization
def send_confirmation_email(email): def send_confirmation_email(email):
confirmation, created = EmailConfirmation.objects.get_or_create(email=email) confirmation, created = EmailConfirmation.objects.get_or_create(email=email)

View File

@@ -4,7 +4,7 @@ from backend.accounts.forms import SignupForm
from .models import EmailConfirmation from .models import EmailConfirmation
from django.shortcuts import get_object_or_404, render from django.shortcuts import get_object_or_404, render
from django.http import HttpResponse 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 from .tasks import create_document_for_organization
class SignUpView(CreateView): class SignUpView(CreateView):
@@ -17,13 +17,10 @@ def confirm_email(request, uuid):
confirmation = get_object_or_404(EmailConfirmation, uuid=uuid) confirmation = get_object_or_404(EmailConfirmation, uuid=uuid)
if confirmation.is_expired(): 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) task = create_document_for_organization.delay(confirmation.email)
send_payment_email(confirmation.email) return render(request, 'accounts/confirmation_success.html',{'email':confirmation.email})
return HttpResponse("Email is confirmed")
def resend_confirmation(request,email): def resend_confirmation(request,email):
if request.method == 'POST': if request.method == 'POST':

View File

@@ -0,0 +1,19 @@
{% extends "base.html" %}
{% block content %}
<div class="container">
<h1>Document Preview</h1>
<h2>Organization: {{ organization.name }}</h2>
<p>{{ created_at }}</p>
<h3>Identified Risks and Controls</h3>
<ul>
{% for risk, controls in risks_with_controls.items %}
<li>
<h2>{{ risk.risk_name }}:</h2>
<p>{{ controls }} </p>
</li>
{% endfor %}
</ul>
</div>
{% endblock %}

View File

@@ -1,11 +1,14 @@
{% extends 'base.html' %} {% extends 'base.html' %}
{% block content %} {% block content %}
<h2>Payment</h2> <div class="out-risk-management">
<p>Click the button below to pay and access your document.</p> <div class="risk-management">
<h2>Payment</h2>
<form method="post"> <p>Click the button below to pay and access your document.</p>
{% csrf_token %} <form method="post">
<button type="submit">Pay</button> {% csrf_token %}
</form> <button class="btn-bl" type="submit">Pay</button>
</form>
</div>
</div>
{% endblock %} {% endblock %}

View File

@@ -12,5 +12,5 @@ urlpatterns = [
path('document/<uuid:document_id>/', v.document, name='document'), path('document/<uuid:document_id>/', v.document, name='document'),
path('preview/<str:name>/', v.template_preview, name='template_preview'), path('preview/<str:name>/', v.template_preview, name='template_preview'),
path("payment/", v.payment_page, name="payment_page"), path("payment/", v.payment_page, name="payment_page"),
path('documentview/<uuid:document_id>/', v.docprew, name='generate_document_view'),
] ]

View File

@@ -1,7 +1,7 @@
from openai import OpenAI from openai import OpenAI
from django.conf import settings from django.conf import settings
from .models import Risk, Control, Document, DocumentRiskControl from .models import Risk, Control, Document, DocumentRiskControl
from django.shortcuts import get_object_or_404 import time
def extract_risk_factors(organization): def extract_risk_factors(organization):
excluded_fields={"name","email"} excluded_fields={"name","email"}
@@ -44,59 +44,101 @@ def get_top_risk(organization):
""" """
response = client.chat.completions.create( response = client.chat.completions.create(
model="gpt-4", model="gpt-4o-mini",
messages=[{"role": "system", "content": prompt}] messages=[{"role": "system", "content": prompt}]
) )
risk_ids = response.choices[0].message.content.strip().split(",") 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()] 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) client = OpenAI(api_key=settings.OPENAI_API_KEY)
all_controls = Control.objects.all() all_controls = Control.objects.all()
control_list = [] control_list = []
risk_factors = extract_risk_factors(organization)
valid_control_ids = {control.id for control in all_controls}
for control in all_controls: for control in all_controls:
control_list.append(f"Control ID: {control.id}, Control Name: {control.name}") control_list.append(f"Control ID: {control.id}, Control Name: {control.name}")
prompt = f""" prompt = f"""
You are a cyber security expert. For the risk '{risk.risk_name}', select 10 relevant controls You are an expert in cybersecurity risk management. Given the risk "{risk.risk_name}" and its associated factors "{risk_factors}",
from the following list and assign a weight (1-10) based on how much they reduce risks. 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.
Available Controls (only respond with control IDs and weights): ### 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} {control_list}
Respond only with control IDs (numbers) and their corresponding weights (1-10).
Format: ### Expected Response Format (STRICTLY FOLLOW THIS FORMAT):
ID: <control_id> Weight: <weight> ```
Example: <control_id> : <weight>
1: 9 <control_id> : <weight>
2: 6
3: 4 ```
### 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( for attempt in range(5):
model="gpt-4", response = client.chat.completions.create(
messages=[{"role": "system", "content": prompt}] 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"): for line in result.split("\n"):
line = line.strip() 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]

View File

@@ -1,13 +1,13 @@
import logging import logging
import yaml
from django.shortcuts import render, redirect , get_object_or_404 from django.shortcuts import render, redirect , get_object_or_404
from .forms import OrganizationForm from .forms import OrganizationForm
from .models import Organization,Document,Risk, DocumentTemplate from .models import Organization,Document,Risk, DocumentTemplate,DocumentRiskControl
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 backend.accounts.utils import send_confirmation_email,send_document_email
from django.contrib.admin.views.decorators import staff_member_required from django.contrib.admin.views.decorators import staff_member_required
# @login_required # @login_required
# def index(request): # def index(request):
# return HttpResponse('<h1>Django</h1><p>Página simples.</p>') # return HttpResponse('<h1>Django</h1><p>Página simples.</p>')
@@ -26,7 +26,6 @@ def signup(request):
if form.is_valid(): if form.is_valid():
form.save() form.save()
send_confirmation_email(form.data['email']) send_confirmation_email(form.data['email'])
return render(request, 'thankyou.html', { return render(request, 'thankyou.html', {
'email': form.data['email'], 'email': form.data['email'],
}) })
@@ -69,3 +68,28 @@ def payment_page(request):
return redirect(document_link) return redirect(document_link)
return render(request, "payment.html", {"email": email}) 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,
})

12
start_services.sh Executable file
View File

@@ -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