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

Dodato full regenerisanje dokumenta,dodat fallback za incomplete document

Browse Source
This commit is contained in:
Amir
2025-08-26 19:32:03 +02:00
parent 959ec958bb
commit 031f955739
8 changed files with 303 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files

199
backend/core/admin.py
View File

@@ -2,7 +2,7 @@ from django.contrib import admin
from .models import Document, DocumentSegment, Organization, Risk, Control, DocumentTemplate, DocumentRiskControl, DemoCode
from django.urls import reverse, path
from django.utils.html import format_html
from .utils import generate_demo_code, get_controls_for_risk, generate_key_findings, generate_recommendations
from .utils import generate_demo_code, get_top_risk, get_controls_for_risk, generate_key_findings, generate_recommendations
from .tables import get_risk_table
from django.shortcuts import render, redirect
from .forms import GenerateCodesForm
@@ -10,6 +10,13 @@ from django.conf import settings
from backend.accounts.utils import send_document_email
from django import forms
from django.contrib.admin.widgets import FilteredSelectMultiple
import logging
from django.contrib import messages
from django.db import transaction
logger = logging.getLogger(__name__)
class DocumentRiskControlInline(admin.TabularInline):
model = DocumentRiskControl
@@ -27,7 +34,7 @@ class DocumentRiskControlInline(admin.TabularInline):
elif obj:
formset.form.base_fields['risk'].queryset = obj.organization.risks.all()
except Exception:
pass
logger.exception("Error building DocumentRiskControlInline formset")
return formset
@@ -65,13 +72,13 @@ class DocumentAdmin(admin.ModelAdmin):
list_filter = ('status', 'created_at')
search_fields = ['organization__name', 'organization__email']
readonly_fields = (
'created_at', 'modified_at',
'regen_controls_action', 'regen_keyfindings_action', 'regen_recommendations_action',
'created_at', 'modified_at', 'regen_note_action',
'regen_document_action', 'regen_top_risks_action', 'regen_controls_action', 'regen_keyfindings_action', 'regen_recommendations_action',
)
fieldsets = (
('Organization & Risks', {
'fields': ('organization', 'organization_risks', 'regen_controls_action')
'fields': ('organization', 'regen_note_action', 'regen_document_action', 'organization_risks', 'regen_top_risks_action')
}),
('Key Findings', {
'fields': ('key_findings', 'regen_keyfindings_action')
@@ -87,18 +94,47 @@ class DocumentAdmin(admin.ModelAdmin):
}),
)
def regen_note_action(self, obj):
return format_html(
'<div class="ai-callout">'
'⏰ <strong>Tip:</strong> Regenerating document can take some time since we depend on AI models to generate content.'
'</div>'
'<br/>'
)
regen_note_action.short_description = ''
def regen_top_risks_action(self,obj):
return format_html(
'<div class="ai-callout ai-callout-info">'
'💡 <strong>Recommended:</strong> after regenerating the Top 10 risks, also update Controls, Key Findings, and Recommendations.'
'</div>'
'<br/>'
'<button type="submit" name="_regen_top_risks" class="button">Regenerate Top 10 Risks using AI</button>'
)
regen_top_risks_action.short_description = ''
def regen_controls_action(self, obj):
return format_html('<button type="submit" name="_regen_controls" class="button">Regenerate Controls using AI</button>')
regen_controls_action.short_description = ''
def regen_keyfindings_action(self, obj):
return format_html('<button type="submit" name="_regen_key_findings" class="button">Regenerate Key Findings using AI</button>')
return format_html('<br/><button type="submit" name="_regen_key_findings" class="button">Regenerate Key Findings using AI</button>')
regen_keyfindings_action.short_description = ''
def regen_recommendations_action(self, obj):
return format_html('<button type="submit" name="_regen_recommendations" class="button">Regenerate Recommendations using AI</button>')
return format_html('<br/><button type="submit" name="_regen_recommendations" class="button">Regenerate Recommendations using AI</button>')
regen_recommendations_action.short_description = ''
def regen_document_action(self, obj):
return format_html(
'<div class="ai-callout">'
'🧹 <strong>Warning:</strong> this will clear the current document (segments, mapped risks/controls, key findings, and recommendations) and regenerate everything as if the document was newly created.'
'</div>'
'<br/>'
'<button type="submit" name="_regen_document" class="button">Regenerate Entire Document</button>'
)
regen_document_action.short_description = ''
def save_model(self, request, obj, form, change):
super().save_model(request, obj, form, change)
org_risks = form.cleaned_data.get('organization_risks')
@@ -111,27 +147,79 @@ class DocumentAdmin(admin.ModelAdmin):
risk_ids = [int(pk) for pk in request.POST.getlist('organization_risks') if pk]
obj.organization.risks.set(Risk.objects.filter(pk__in=risk_ids))
except Exception:
pass
logger.exception("Failed to apply posted organization_risks")
def _clear_segments(self, obj, startswith=None, exact=None):
try:
if startswith:
for s in startswith:
obj.segments.filter(content__startswith=s).delete()
if exact:
obj.segments.filter(content__in=exact).delete()
except Exception:
logger.exception("Failed to clear segments for document %s", getattr(obj, 'pk', None))
def _risk_content(self, risks):
return "\n\n".join([
f"Risk: {risk.risk_id} - {risk.risk_name} \n"
f"Category: {risk.category}\n"
f"Primary Impact: {risk.primary_impact} \n"
f"Secondary Impact: {risk.secondary_impact}\n"
f"Tertiary Impact: {risk.tretiary_impact} \n"
f"Detection Difficulty: {risk.detection_difficulty} \n"
f"Recovery Complexity: {risk.recovery_complexity} \n"
f"Business Impact Severity: {risk.businnes_impact_severity}\n"
for risk in risks
])
def _add_identified_risks(self, obj, risks):
if not risks:
return
self._clear_segments(obj, startswith=["Identified Risks"], exact=["Top 10 Risks Identified"])
obj.add_segment('h1', "Top 10 Risks Identified")
obj.add_segment('body', f"Identified Risks: \n\n{self._risk_content(risks)}")
def _clear_document_mappings(self, obj, clear_org_risks=True):
try:
obj.segments.all().delete()
obj.documentriskcontrol_set.all().delete()
obj.key_findings = ''
obj.recomendations = ''
obj.status = Document.STATUS_WAITING
obj.save(update_fields=['key_findings', 'recomendations', 'status', 'modified_at'])
if clear_org_risks and getattr(obj, 'organization', None):
obj.organization.risks.clear()
except Exception:
logger.exception("Failed to clear document mappings for document %s", getattr(obj, 'pk', None))
def _regen_pipeline(self, obj):
ok = True
if not self._regen_top_risks(obj):
ok = False
else:
try:
self._regen_controls(obj)
except Exception:
logger.exception("_regen_controls failed")
ok = False
if not self._regen_key_findings(obj):
ok = False
if not self._regen_recommendations(obj):
ok = False
return ok
def _regen_top_risks(self, obj):
top_risk_ids = get_top_risk(obj.organization)
top_risks = Risk.objects.filter(risk_id__in=top_risk_ids)
obj.organization.risks.set(top_risks)
self._add_identified_risks(obj, top_risks)
return True
def _regen_controls(self, obj):
obj.segments.filter(content__startswith="Identified Risks").delete()
obj.segments.filter(content__startswith="Mitigation Controls").delete()
obj.segments.filter(content__in=["Top 10 Risks Identified", "Regenerated Controls"]).delete()
self._clear_segments(obj, startswith=["Identified Risks", "Mitigation Controls"], exact=["Top 10 Risks Identified", "Regenerated Controls"])
obj.documentriskcontrol_set.all().delete()
top_risks = list(obj.organization.risks.all())
obj.add_segment('h1', "Top 10 Risks Identified")
risk_content = "\n\n".join([
f"Risk: {r.risk_id} - {r.risk_name} \n"
f"Category: {r.category}\n"
f"Primary Impact: {r.primary_impact} \n"
f"Secondary Impact: {r.secondary_impact}\n"
f"Tertiary Impact: {r.tretiary_impact} \n"
f"Detection Difficulty: {r.detection_difficulty} \n"
f"Recovery Complexity: {r.recovery_complexity} \n"
f"Business Impact Severity: {r.businnes_impact_severity}\n"
for r in top_risks
])
obj.add_segment('body', f"Identified Risks: \n\n{risk_content}")
self._add_identified_risks(obj, top_risks)
controls_content = "Mitigation Controls:\n\n"
for risk in top_risks:
controls_content += f"Risk: {risk.risk_id} - {risk.risk_name}\n"
@@ -170,27 +258,55 @@ class DocumentAdmin(admin.ModelAdmin):
return False
def changeform_view(self, request, object_id=None, form_url='', extra_context=None):
if request.method == 'POST' and any(k in request.POST for k in ("_regen_controls", "_regen_key_findings", "_regen_recommendations")):
if request.method == 'POST' and any(k in request.POST for k in ("_regen_controls", "_regen_key_findings", "_regen_recommendations", "_regen_top_risks", "_regen_document")):
obj = self.get_object(request, object_id)
if obj is None:
return super().changeform_view(request, object_id, form_url, extra_context)
try:
self._apply_post_org_risks(request, obj)
if "_regen_controls" in request.POST:
if "_regen_top_risks" in request.POST:
if not obj.organization_id:
self.message_user(request, "Please select an organization first.", level=messages.WARNING)
else:
if self._regen_top_risks(obj):
self.message_user(request, "Top risks regenerated and risk segment updated.")
else:
self.message_user(request, "Top risks could not be generated.", level=messages.WARNING)
elif "_regen_controls" in request.POST:
self._regen_controls(obj)
self.message_user(request, "Risks and controls regenerated successfully.")
elif "_regen_key_findings" in request.POST:
if self._regen_key_findings(obj):
self.message_user(request, "Key Findings regenerated.")
else:
self.message_user(request, "Key Findings could not be generated.", level='warning')
self.message_user(request, "Key Findings could not be generated.", level=messages.WARNING)
elif "_regen_recommendations" in request.POST:
if self._regen_recommendations(obj):
self.message_user(request, "Recommendations regenerated.")
else:
self.message_user(request, "Recommendations could not be generated.", level='warning')
self.message_user(request, "Recommendations could not be generated.", level=messages.WARNING)
elif "_regen_document" in request.POST:
if not obj.organization_id:
self.message_user(request, "Please select an organization first.", level=messages.WARNING)
else:
try:
with transaction.atomic():
self._clear_document_mappings(obj, clear_org_risks=True)
regen_ok = self._regen_pipeline(obj)
try:
self.log_change(request, obj, "Full document regeneration triggered")
except Exception:
logger.exception("Failed to log admin change for full regen")
if regen_ok:
self.message_user(request, "Document fully regenerated.")
else:
self.message_user(request, "Document regeneration finished with warnings or missing outputs.", level=messages.WARNING)
except Exception as e:
logger.exception("Full regeneration failed")
self.message_user(request, f"Full regeneration failed: {e}", level=messages.ERROR)
except Exception as e:
self.message_user(request, f"Action failed: {e}", level='error')
logger.exception("changeform_view action failed")
self.message_user(request, f"Action failed: {e}", level=messages.ERROR)
return redirect(reverse('admin:core_document_change', args=[obj.pk]))
return super().changeform_view(request, object_id, form_url, extra_context)
@@ -209,32 +325,17 @@ class DocumentAdmin(admin.ModelAdmin):
obj.save(update_fields=['status', 'modified_at'])
self.message_user(request, "Document sent and marked as done.")
except Exception as e:
self.message_user(request, f"Failed to send document: {e}", level='error')
logger.exception("Failed to send document email")
self.message_user(request, f"Failed to send document: {e}", level=messages.ERROR)
return redirect(reverse('admin:core_document_change', args=[obj.pk]))
return super().response_change(request, obj)
def _refresh_segments_from_current_mappings(self, obj):
obj.segments.filter(content__startswith="Identified Risks").delete()
obj.segments.filter(content__startswith="Mitigation Controls").delete()
obj.segments.filter(content__in=["Top 10 Risks Identified", "Regenerated Controls"]).delete()
self._clear_segments(obj, startswith=["Identified Risks", "Mitigation Controls"], exact=["Top 10 Risks Identified", "Regenerated Controls"])
top_risks = list(obj.organization.risks.all())
if top_risks:
obj.add_segment('h1', "Top 10 Risks Identified")
risk_content = "\n\n".join([
f"Risk: {r.risk_id} - {r.risk_name} \n"
f"Category: {r.category}\n"
f"Primary Impact: {r.primary_impact} \n"
f"Secondary Impact: {r.secondary_impact}\n"
f"Tertiary Impact: {r.tretiary_impact} \n"
f"Detection Difficulty: {r.detection_difficulty} \n"
f"Recovery Complexity: {r.recovery_complexity} \n"
f"Business Impact Severity: {r.businnes_impact_severity}\n"
for r in top_risks
])
obj.add_segment('body', f"Identified Risks: \n\n{risk_content}")
self._add_identified_risks(obj, top_risks)
from collections import defaultdict
controls_by_risk = defaultdict(list)
@@ -260,7 +361,7 @@ class DocumentAdmin(admin.ModelAdmin):
try:
self._refresh_segments_from_current_mappings(obj)
except Exception:
pass
logger.exception("Failed to refresh segments from current mappings for document %s", getattr(obj, 'pk', None))
class DocumentTemplateAdmin(admin.ModelAdmin):
list_display = ['name', 'created_at', 'updated_at', 'preview_button']