Added expert review/edit document
This commit is contained in:
@@ -1,23 +1,266 @@
|
||||
from django.contrib import admin
|
||||
from .models import Document, DocumentSegment, Organization, Risk, Control, DocumentTemplate, DocumentRiskControl, DemoCode
|
||||
from django.urls import reverse
|
||||
from django.urls import reverse, path
|
||||
from django.utils.html import format_html
|
||||
from .utils import generate_demo_code
|
||||
from django.urls import path
|
||||
from .utils import generate_demo_code, 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
|
||||
from django.conf import settings
|
||||
from backend.accounts.utils import send_document_email
|
||||
from django import forms
|
||||
from django.contrib.admin.widgets import FilteredSelectMultiple
|
||||
|
||||
class DocumentRiskControlInline(admin.TabularInline):
|
||||
model = DocumentRiskControl
|
||||
extra = 2
|
||||
max_num = 10
|
||||
can_delete = False
|
||||
fields = ('risk', 'control', 'weight', 'likelihood')
|
||||
|
||||
def get_formset(self, request, obj=None, **kwargs):
|
||||
formset = super().get_formset(request, obj, **kwargs)
|
||||
try:
|
||||
if request.method == 'POST' and 'organization_risks' in request.POST:
|
||||
risk_ids = request.POST.getlist('organization_risks')
|
||||
formset.form.base_fields['risk'].queryset = Risk.objects.filter(pk__in=risk_ids)
|
||||
elif obj:
|
||||
formset.form.base_fields['risk'].queryset = obj.organization.risks.all()
|
||||
except Exception:
|
||||
pass
|
||||
return formset
|
||||
|
||||
|
||||
class DocumentAdminForm(forms.ModelForm):
|
||||
organization_risks = forms.ModelMultipleChoiceField(
|
||||
queryset=Risk.objects.all(),
|
||||
required=False,
|
||||
widget=FilteredSelectMultiple(verbose_name="Risks", is_stacked=False),
|
||||
help_text="Edit the AI-selected risks for this organization."
|
||||
)
|
||||
|
||||
class Meta:
|
||||
model = Document
|
||||
fields = ['organization', 'status', 'key_findings', 'recomendations']
|
||||
|
||||
class Media:
|
||||
css = { 'all': ('admin/css/widgets.css',) }
|
||||
js = (
|
||||
'admin/js/SelectBox.js',
|
||||
'admin/js/SelectFilter2.js',
|
||||
)
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
if self.instance and getattr(self.instance, 'organization_id', None):
|
||||
self.fields['organization_risks'].initial = self.instance.organization.risks.all()
|
||||
|
||||
|
||||
class DocumentSegmentInline(admin.StackedInline):
|
||||
model = DocumentSegment
|
||||
extra = 1
|
||||
ordering = ['order']
|
||||
fields = ('segment_type', 'content', 'order')
|
||||
|
||||
class DocumentAdmin(admin.ModelAdmin):
|
||||
inlines = [DocumentSegmentInline]
|
||||
list_display = ('organization', 'created_at', 'modified_at')
|
||||
search_fields = ['organization__name']
|
||||
readonly_fields = ('created_at', 'modified_at')
|
||||
change_form_template = "admin/core/document/change_form.html"
|
||||
|
||||
form = DocumentAdminForm
|
||||
inlines = [DocumentRiskControlInline]
|
||||
list_display = ('organization', 'status', 'created_at', 'modified_at', 'review_link')
|
||||
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',
|
||||
)
|
||||
|
||||
fieldsets = (
|
||||
('Organization & Risks', {
|
||||
'fields': ('organization', 'organization_risks', 'regen_controls_action')
|
||||
}),
|
||||
('Key Findings', {
|
||||
'fields': ('key_findings', 'regen_keyfindings_action')
|
||||
}),
|
||||
('Recommendations', {
|
||||
'fields': ('recomendations', 'regen_recommendations_action')
|
||||
}),
|
||||
('Status', {
|
||||
'fields': ('status',)
|
||||
}),
|
||||
('Timestamps', {
|
||||
'fields': ('created_at', 'modified_at')
|
||||
}),
|
||||
)
|
||||
|
||||
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>')
|
||||
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>')
|
||||
regen_recommendations_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')
|
||||
if org_risks is not None and obj.organization_id:
|
||||
obj.organization.risks.set(org_risks)
|
||||
|
||||
def _apply_post_org_risks(self, request, obj):
|
||||
try:
|
||||
if 'organization_risks' in request.POST and obj.organization_id:
|
||||
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
|
||||
|
||||
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()
|
||||
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}")
|
||||
controls_content = "Mitigation Controls:\n\n"
|
||||
for risk in top_risks:
|
||||
controls_content += f"Risk: {risk.risk_id} - {risk.risk_name}\n"
|
||||
selected_controls = get_controls_for_risk(risk, organization=obj.organization)
|
||||
for control_id, weight, likelihood in selected_controls:
|
||||
control = Control.objects.filter(id=control_id).first()
|
||||
if control:
|
||||
DocumentRiskControl.objects.create(
|
||||
document=obj,
|
||||
risk=risk,
|
||||
control=control,
|
||||
weight=weight,
|
||||
likelihood=likelihood
|
||||
)
|
||||
label = f"{control.subcategory} - {control.function or ''}".rstrip(" -")
|
||||
controls_content += f" - Control: {label} (Impact Weight: {weight}/10) (Likelihood: {likelihood}/10)\n"
|
||||
controls_content += "\n"
|
||||
obj.add_segment('body', controls_content)
|
||||
|
||||
def _regen_key_findings(self, obj):
|
||||
risks_top3 = get_risk_table(obj)[:3]
|
||||
key_findings = generate_key_findings(obj, risks_top3)
|
||||
if key_findings:
|
||||
obj.key_findings = key_findings
|
||||
obj.save(update_fields=['key_findings', 'modified_at'])
|
||||
return True
|
||||
return False
|
||||
|
||||
def _regen_recommendations(self, obj):
|
||||
risks_top10 = get_risk_table(obj)[:10]
|
||||
recommendations = generate_recommendations(risks_top10, obj.organization)
|
||||
if recommendations:
|
||||
obj.recomendations = recommendations
|
||||
obj.save(update_fields=['recomendations', 'modified_at'])
|
||||
return True
|
||||
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")):
|
||||
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:
|
||||
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')
|
||||
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')
|
||||
except Exception as e:
|
||||
self.message_user(request, f"Action failed: {e}", level='error')
|
||||
return redirect(reverse('admin:core_document_change', args=[obj.pk]))
|
||||
return super().changeform_view(request, object_id, form_url, extra_context)
|
||||
|
||||
def review_link(self, obj):
|
||||
url = reverse('admin:core_document_change', args=[obj.pk])
|
||||
label = 'Review / Edit'
|
||||
return format_html('<a class="button" href="{}">{}</a>', url, label)
|
||||
review_link.short_description = 'Action'
|
||||
|
||||
def response_change(self, request, obj):
|
||||
if "_save_send" in request.POST:
|
||||
try:
|
||||
url = f"{settings.SITE_DOMAIN}/pdf/{obj.id}/"
|
||||
send_document_email(obj.organization.email, url, obj)
|
||||
obj.status = Document.STATUS_DONE
|
||||
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')
|
||||
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()
|
||||
|
||||
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}")
|
||||
|
||||
from collections import defaultdict
|
||||
controls_by_risk = defaultdict(list)
|
||||
for drc in obj.documentriskcontrol_set.select_related('risk', 'control').all():
|
||||
controls_by_risk[drc.risk_id].append(drc)
|
||||
|
||||
controls_content = "Mitigation Controls:\n\n"
|
||||
for risk in top_risks:
|
||||
controls_content += f"Risk: {risk.risk_id} - {risk.risk_name}\n"
|
||||
rows = controls_by_risk.get(risk.pk, [])
|
||||
rows.sort(key=lambda x: (x.control.subcategory or '', x.control.function or ''))
|
||||
for drc in rows:
|
||||
control = drc.control
|
||||
if control:
|
||||
label = f"{control.subcategory} - {control.function or ''}".rstrip(" -")
|
||||
controls_content += f" - Control: {label} (Impact Weight: {drc.weight}/10) (Likelihood: {drc.likelihood}/10)\n"
|
||||
controls_content += "\n"
|
||||
obj.add_segment('body', controls_content)
|
||||
|
||||
def save_related(self, request, form, formsets, change):
|
||||
super().save_related(request, form, formsets, change)
|
||||
obj = form.instance
|
||||
try:
|
||||
self._refresh_segments_from_current_mappings(obj)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
class DocumentTemplateAdmin(admin.ModelAdmin):
|
||||
list_display = ['name', 'created_at', 'updated_at', 'preview_button']
|
||||
|
||||
Reference in New Issue
Block a user