old-riskletpy/backend/core/utils.py

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]