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

Resen konflikt

Browse Source
This commit is contained in:
Amir
2025-02-21 00:32:27 +01:00
parent 5ed323140f
commit 3175f9e7cb
6 changed files with 134 additions and 148 deletions
Download Patch File Download Diff File Expand all files Collapse all files

152
backend/core/utils.py
View File

@@ -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):
```
<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.**
"""
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):
<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.**
"""
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 []