senaduka/old-riskletpy
1
0
Fork 0
Code Issues 5 Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Changed to gpt-4o-mini, Added organisation info

Browse Source
This commit is contained in:
Amir
2025-02-17 21:40:08 +01:00
parent 65716ff842
commit 7151cd0c4d
2 changed files with 80 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
backend/accounts/tasks.py
View File

@@ -34,7 +34,7 @@ def create_document_for_organization(confirmation_email):
for risk in top_risks:
controls_content += f"Risk: {risk.risk_id} - {risk.risk_name}\n"
selected_controls = get_controls_for_risk(risk)
selected_controls = get_controls_for_risk(risk ,organization=organization)
for control_id, weight in selected_controls:
control = Control.objects.filter(id=control_id).first()

114
backend/core/utils.py
View File

@@ -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
from .models import Risk, Control
import time
def extract_risk_factors(organization):
excluded_fields={"name","email"}
@@ -26,6 +26,11 @@ 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)
@@ -44,59 +49,98 @@ 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: <control_id> Weight: <weight>
Example:
1: 9
2: 6
3: 4
### Expected Response Format (STRICTLY FOLLOW THIS FORMAT):
```
<control_id> : <weight>
<control_id> : <weight>
```
### 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("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()
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))
control_id = int(control_id_str)
weight = int(weight_str)
print(f"ID: {control_id}, Weight: {weight}")
control_id_str = ''.join(filter(str.isdigit, control_id_str))
weight_str = ''.join(filter(str.isdigit, weight_str))
control = Control.objects.filter(id=control_id).first()
if control:
selected_controls.append((control_id, weight))
return selected_controls[:10]
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 []