from celery import shared_task from backend.core.models import Organization, Document, Risk, Control, DocumentRiskControl from backend.core.utils import get_top_risk, get_controls_for_risk, generate_key_findings, generate_recommendations from django.shortcuts import get_object_or_404, render from .utils import send_payment_email, send_document_to_reviewer from backend.core.tables import get_risk_table import logging logger = logging.getLogger(__name__) @shared_task def create_document_for_organization(confirmation_email): is_incomplete = False organization = get_object_or_404(Organization, email=confirmation_email) top_risks_with_explanation = get_top_risk(organization) top_risk_ids = [r['risk_id'] for r in top_risks_with_explanation] if len(top_risk_ids) != 10: is_incomplete = True top_risks_with_explanation = get_top_risk(organization) top_risk_ids = [r['risk_id'] for r in top_risks_with_explanation] top_risks = Risk.objects.filter(risk_id__in=top_risk_ids) organization.risks.set(top_risks) explanation_map = {r['risk_id']: r['explanation'] for r in top_risks_with_explanation} document = Document.objects.create(organization=organization) document.risk_explanations = explanation_map document.add_segment('h1', "Top 10 Risks Identified") risk_content = "\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" f"Explanation: {explanation_map.get(risk.risk_id, '')}\n" for risk in top_risks ]) document.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=organization) if len(selected_controls) < 10: is_incomplete = True selected_controls = get_controls_for_risk(risk, organization=organization) for control_id, weight, likelihood in selected_controls: control = Control.objects.filter(id=control_id).first() if control: DocumentRiskControl.objects.create( document=document, 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" document.add_segment('body', controls_content) risks_by_weight = get_risk_table(document)[:3] key_findings = generate_key_findings(document, risks_by_weight) if key_findings: document.key_findings = key_findings document.save() risk_for_recomendation = get_risk_table(document)[:10] recommendations = generate_recommendations(risk_for_recomendation, organization) if recommendations: document.recomendations = recommendations document.save() if is_incomplete: logger.info("Marking document %s as INCOMPLETE (is_incomplete=%s)", document.id, is_incomplete) document.status = Document.STATUS_INCOMPLETE document.save(update_fields=['status', 'modified_at']) send_payment_email(confirmation_email)