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

Zamenio na while petlju dok ne popuni svi 10 kontrola za rizik

Browse Source
This commit is contained in:
Amir
2025-02-18 21:49:07 +01:00
parent 51a6701794
commit 5a0c4df816
1 changed files with 79 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files

153
backend/core/utils.py
View File

@@ -63,14 +63,17 @@ 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 = []
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 organization details "{organization_details}",
@@ -85,95 +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(10):
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:
continue
selected_controls.append((control_id, weight))
control_ids_seen.add(control_id)
else:
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 valid and len(selected_controls) == 10:
return selected_controls
print(f"Recived {len(selected_controls)} controls. Retrying for missing ones...\n")
remaining_controls = valid_control_ids - control_ids_seen
missing_count = 10 - len(selected_controls)
if missing_count > 0:
remaining_controls_list = [f"Control ID:{cid}, Control Name: {control_map[cid]}" for cid in remaining_controls]
prompt = f"""
You are an expert in cybersecurity risk management. Previously, you selected {len(selected_controls)} controls for the risk "{risk.risk_name}"
and its associated organization details "{organization_details}".
Now, your task is to select **exactly {missing_count} additional unique controls** from the remaining list that best mitigate this risk.
Each selected 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 with the previously selected controls).
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.**
### Remaining Available Controls:
{remaining_controls_list}
### Expected Response Format (STRICTLY FOLLOW THIS FORMAT):
<control_id> : <weight>
<control_id> : <weight>
### Example Correct Response (NO DUPLICATES):
23 : 7,
45 : 3,
**If you provide duplicate control IDs or controls outside the remaining list, your response will be rejected.**
**Follow the response format exactly. Any deviation will be considered invalid.**
"""
time.sleep(2)
continue
print("Failed to get a valid response after multiple attempts.")
return []
if not remaining_controls:
break
return selected_controls if len(selected_controls) == 10 else []