Added graph and table to document

This commit is contained in:
2025-03-27 23:57:31 +01:00
parent 5a7a89d93c
commit fc29671331
16 changed files with 687 additions and 65 deletions

View File

@@ -5,6 +5,14 @@ from weasyprint import HTML
from django.http import HttpResponse
from PIL import Image
import io
import base64
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from django.contrib.staticfiles.finders import find
import matplotlib.image as mpimg
def extract_organization_details(organization):
excluded_fields = {"name", "email"}
@@ -59,7 +67,6 @@ def get_top_risk(organization):
)
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()]
@@ -80,22 +87,28 @@ def get_controls_for_risk(risk, organization):
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 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.
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** (1 = low impact, 10 = high impact).
- A likelihood score between **1 and 10** (1 = rare occurrence, 10 = highly likely).
### Rules:
1. **Each control ID must be unique** (no duplicates).
2. **Only return control IDs and weights** in the exact format below.
2. **Only return control IDs, weights, and likelihood scores** 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.**
4. **Likelihood scores must be between 1 and 10** (1 = rare occurrence, 10 = highly likely).
5. **Do NOT add explanations, descriptions, or extra text.**
6. **Ensure that control IDs are randomly distributed and diverse across different categories.**
### Available Controls:
{control_list}
### Expected Response Format (STRICTLY FOLLOW THIS FORMAT):
<control_id> : <weight>
<control_id> : <weight>
<control_id> : <weight> : <likelihood>
<control_id> : <weight> : <likelihood>
### Example Correct Response (NO DUPLICATES):
12 : 8
45 : 7
12 : 8 : 90
45 : 7 : 60
⚠️ **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.**
"""
@@ -108,22 +121,23 @@ def get_controls_for_risk(risk, organization):
for line in result.split("\n"):
line = line.strip()
parts = line.split(":")
if len(parts) == 2:
if len(parts) == 3:
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)}")
likelihood_str = parts[2].strip().replace("Likelihood:", "").replace("likelihood:", "").strip()
control_id_str = ''.join(filter(str.isdigit, control_id_str))
weight_str = ''.join(filter(str.isdigit, weight_str))
likelihood_str = ''.join(filter(str.isdigit, likelihood_str))
if control_id_str and weight_str:
if control_id_str and weight_str and likelihood_str:
try:
control_id = int(control_id_str)
weight = int(weight_str)
likelihood = int(likelihood_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))
if control_id in valid_control_ids and 1 <= weight <= 10 and 1 <= likelihood <= 10 and control_id not in control_ids_seen:
selected_controls.append((control_id, weight, likelihood))
control_ids_seen.add(control_id)
except ValueError:
continue
@@ -137,23 +151,30 @@ def get_controls_for_risk(risk, organization):
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.
You are an expert in cybersecurity risk management. Given the risk "{risk.risk_name}" and the organization's 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.
- A likelihood score between **1 and 10** (1 = rare occurrence, 10 = highly likely).
### Rules:
1. **Each control ID must be unique** (no duplicates).
2. **Only return control IDs and weights** in the exact format below.
2. **Only return control IDs, weights, and likelihood scores** 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.**
4. **Likelihood scores must be between 1 and 10** (1 = rare occurrence, 10 = highly likely).
5. **Do NOT add explanations, descriptions, or extra text.**
6. **Ensure that control IDs are diverse and well-distributed across different categories.**
### Available Controls:
{remaining_controls_list}
### Expected Response Format (STRICTLY FOLLOW THIS FORMAT):
<control_id> : <weight>
<control_id> : <weight>
<control_id> : <weight> : <likelihood>
<control_id> : <weight> : <likelihood>
### Example Correct Response (NO DUPLICATES):
12 : 8
45 : 7
12 : 8 : 85
45 : 7 : 60
⚠️ **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.**
"""
@@ -162,22 +183,24 @@ def get_controls_for_risk(risk, organization):
for line in result.split("\n"):
line = line.strip()
parts = line.split(":")
if len(parts) == 2:
if len(parts) == 3:
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)}")
likelihood_str = parts[2].strip().replace("Likelihood:", "").replace("likelihood:", "").strip()
control_id_str = ''.join(filter(str.isdigit, control_id_str))
weight_str = ''.join(filter(str.isdigit, weight_str))
likelihood_str = ''.join(filter(str.isdigit, likelihood_str))
if control_id_str and weight_str:
if control_id_str and weight_str and likelihood_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))
likelihood = int(likelihood_str)
if control_id in valid_control_ids and 1 <= weight <= 10 and 1 <= likelihood <= 10 and control_id not in control_ids_seen:
selected_controls.append((control_id, weight, likelihood))
control_ids_seen.add(control_id)
except ValueError:
continue
@@ -207,4 +230,68 @@ def generate_first_page_image(document):
images[0].save(img_io, format="JPEG", quality=90)
img_io.seek(0)
return img_io
return img_io
def calculate_aggregate_weight(controls):
total_weight = sum(control['weight']for control in controls)
return total_weight
def calculate_aggregate_likelihood(controls):
total_likelihood = sum(control['likelihood'] for control in controls)
return total_likelihood
def map_weight_to_impact_likelihood(total_weight, total_likelihood, max_weight):
normalized_weight = total_weight / max_weight
impact = min(10.0, max(1.0, normalized_weight * 10.0))
likelihood = min(10.0, max(1.0, total_likelihood / 10.0))
return impact, likelihood
def generate_risk_graph(risks_with_controls):
impacts = [risk['impact'] for risk in risks_with_controls]
likelihoods = [risk['likelihood'] for risk in risks_with_controls]
risk_ids = [risk['risk']['id'] for risk in risks_with_controls]
bg_img_path = find('img/graph_matrix (3).png')
bg_img = mpimg.imread(bg_img_path)
fig, ax = plt.subplots(figsize=(10, 8))
ax.imshow(bg_img, extent=[0, 11.2, 0, 11.2], aspect='auto')
scatter = ax.scatter(
likelihoods, impacts,
c="blue", edgecolors="white", s=500, alpha=0.9
)
for i, risk_id in enumerate(risk_ids):
ax.annotate(
str(risk_id),
(likelihoods[i], impacts[i]),
color="white",
fontsize=12,
ha="center",
va="center",
weight="bold",
)
ax.set_xticks([])
ax.set_yticks([])
ax.set_xticklabels([])
ax.set_yticklabels([])
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.spines['left'].set_visible(False)
ax.spines['bottom'].set_visible(False)
buffer = io.BytesIO()
plt.savefig(buffer, format="png", transparent=True, bbox_inches='tight', pad_inches=0)
buffer.seek(0)
image_png = buffer.getvalue()
buffer.close()
plt.close()
return base64.b64encode(image_png).decode("utf-8")