| 1 | import os |
| 2 | from tempfile import NamedTemporaryFile |
| 3 | from typing import Dict |
| 4 | from typing import Optional |
| 5 | |
| 6 | from docxtpl import DocxTemplate |
| 7 | from fastapi.responses import FileResponse |
| 8 | |
| 9 | from app.data_store.data_store_operations import download_data_store |
| 10 | |
| 11 | |
| 12 | async def download_template(template_name: str) -> bytes: |
| 13 | """Retrieve the template file content from the data store.""" |
| 14 | return await download_data_store(bucket_name="copilot-case-report-templates", object_name=template_name) |
| 15 | |
| 16 | |
| 17 | def create_case_context(case) -> Dict[str, Dict[str, str]]: |
| 18 | """Prepare the context for the Jinja template.""" |
| 19 | return { |
| 20 | "case": { |
| 21 | "name": case.case_name, |
| 22 | "description": case.case_description, |
| 23 | "assigned_to": case.assigned_to, |
| 24 | "case_creation_time": case.case_creation_time, |
| 25 | "case_closed_time": case.case_closed_time, |
| 26 | "id": case.id, |
| 27 | "alerts": [ |
| 28 | { |
| 29 | "alert_name": alert.alert.alert_name, |
| 30 | "alert_description": alert.alert.alert_description, |
| 31 | "status": alert.alert.status, |
| 32 | "time_closed": alert.alert.time_closed, |
| 33 | "tags": [tag.tag.tag for tag in alert.alert.tags], |
| 34 | "assets": [ |
| 35 | { |
| 36 | "asset_name": asset.asset_name, |
| 37 | "agent_id": asset.agent_id, |
| 38 | } |
| 39 | for asset in alert.alert.assets |
| 40 | ], |
| 41 | "comments": [ |
| 42 | { |
| 43 | "comment": comment.comment, |
| 44 | "user_name": comment.user_name, |
| 45 | "created_at": comment.created_at, |
| 46 | } |
| 47 | for comment in alert.alert.comments |
| 48 | ], |
| 49 | "context": { |
| 50 | "source": alert.alert.assets[0].alert_context.source |
| 51 | if alert.alert.assets and alert.alert.assets[0].alert_context |
| 52 | else None, |
| 53 | "context": alert.alert.assets[0].alert_context.context |
| 54 | if alert.alert.assets and alert.alert.assets[0].alert_context |
| 55 | else None, |
| 56 | } |
| 57 | if alert.alert.assets |
| 58 | else None, |
| 59 | "iocs": [ |
| 60 | { |
| 61 | "ioc_value": ioc.ioc.value, |
| 62 | "ioc_type": ioc.ioc.type, |
| 63 | "ioc_description": ioc.ioc.description, |
| 64 | } |
| 65 | for ioc in alert.alert.iocs |
| 66 | ], |
| 67 | } |
| 68 | for alert in case.alerts |
| 69 | ], |
| 70 | }, |
| 71 | } |
| 72 | |
| 73 | |
| 74 | def save_template_to_tempfile(template_file_content: bytes) -> str: |
| 75 | """Save the template content to a temporary file.""" |
| 76 | with NamedTemporaryFile(delete=False, suffix=".docx") as tmp_template: |
| 77 | tmp_template.write(template_file_content) |
| 78 | return tmp_template.name |
| 79 | |
| 80 | |
| 81 | def render_document_with_context(template_path: str, context: Dict[str, Dict[str, str]]) -> str: |
| 82 | """Load and render the document template with the given context.""" |
| 83 | doc = DocxTemplate(template_path) |
| 84 | doc.render(context) |
| 85 | with NamedTemporaryFile(delete=False, suffix=".docx") as tmp: |
| 86 | doc.save(tmp.name) |
| 87 | return tmp.name |
| 88 | |
| 89 | |
| 90 | def create_file_response(file_path: str, file_name: Optional[str] = "case_report.docx") -> FileResponse: |
| 91 | """Create a FileResponse object for the rendered document.""" |
| 92 | return FileResponse( |
| 93 | file_path, |
| 94 | filename=file_name, |
| 95 | media_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document", |
| 96 | ) |
| 97 | |
| 98 | |
| 99 | def cleanup_temp_files(file_paths: list): |
| 100 | """Clean up the temporary files.""" |
| 101 | for file_path in file_paths: |
| 102 | if os.path.exists(file_path): |
| 103 | os.remove(file_path) |