103 lines
3.3 KiB
Python
103 lines
3.3 KiB
Python
from openai import OpenAI
|
|
from django.conf import settings
|
|
from .models import Risk, Control, Document, DocumentRiskControl
|
|
from django.shortcuts import get_object_or_404
|
|
|
|
def extract_risk_factors(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
|
|
return risk_data
|
|
|
|
def get_top_risk(organization):
|
|
client = OpenAI(api_key=settings.OPENAI_API_KEY)
|
|
|
|
all_risks = Risk.objects.all()
|
|
|
|
risk_list = []
|
|
for risk in all_risks:
|
|
risk_list.append(f"""
|
|
Risk ID: {risk.risk_id}
|
|
Category: {risk.category}
|
|
Name: {risk.risk_name}
|
|
Primary Impact: {risk.primary_impact}
|
|
""")
|
|
|
|
risk_factors = extract_risk_factors(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}
|
|
|
|
List of Risks:
|
|
{risk_list}
|
|
|
|
Provide only the 10 most critical risk IDs in a simple comma-separated format, e.g "1,3,7,12,..."
|
|
"""
|
|
|
|
response = client.chat.completions.create(
|
|
model="gpt-4",
|
|
messages=[{"role": "system", "content": prompt}]
|
|
)
|
|
|
|
risk_ids = response.choices[0].message.content.strip().split(",")
|
|
|
|
return [int(risk_id) for risk_id in risk_ids if risk_id.isdigit()]
|
|
|
|
def get_controls_for_risk(risk):
|
|
client = OpenAI(api_key=settings.OPENAI_API_KEY)
|
|
all_controls = Control.objects.all()
|
|
control_list = []
|
|
|
|
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):
|
|
{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
|
|
"""
|
|
|
|
response = client.chat.completions.create(
|
|
model="gpt-4",
|
|
messages=[{"role": "system", "content": prompt}]
|
|
)
|
|
|
|
result = response.choices[0].message.content.strip()
|
|
selected_controls = []
|
|
|
|
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()
|
|
|
|
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 = Control.objects.filter(id=control_id).first()
|
|
if control:
|
|
selected_controls.append((control_id, weight))
|
|
return selected_controls[:10]
|