Added graph and table to document
This commit is contained in:
@@ -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")
|
||||
Reference in New Issue
Block a user