| 1 | import os |
| 2 | import platform |
| 3 | from tempfile import NamedTemporaryFile |
| 4 | from typing import Dict |
| 5 | from typing import Optional |
| 6 | |
| 7 | import pdfkit |
| 8 | from fastapi.responses import FileResponse |
| 9 | from jinja2 import Environment |
| 10 | from jinja2 import FileSystemLoader |
| 11 | |
| 12 | from app.data_store.data_store_operations import download_data_store |
| 13 | |
| 14 | |
| 15 | async def download_template_pdf(template_name: str) -> str: |
| 16 | """Retrieve the template file content from the data store and save it to a temporary file.""" |
| 17 | template_content = await download_data_store(bucket_name="copilot-case-report-templates", object_name=template_name) |
| 18 | with NamedTemporaryFile(delete=False, suffix=".html") as tmp_template: |
| 19 | tmp_template.write(template_content) |
| 20 | return tmp_template.name |
| 21 | |
| 22 | |
| 23 | def create_case_context_pdf(case) -> Dict[str, Dict[str, str]]: |
| 24 | """Prepare the context for the Jinja template.""" |
| 25 | return { |
| 26 | "case": { |
| 27 | "name": case.case_name, |
| 28 | "description": case.case_description, |
| 29 | "assigned_to": case.assigned_to, |
| 30 | "case_creation_time": case.case_creation_time, |
| 31 | "case_closed_time": case.case_closed_time, |
| 32 | "id": case.id, |
| 33 | "alerts": [ |
| 34 | { |
| 35 | "alert_name": alert.alert.alert_name, |
| 36 | "alert_description": alert.alert.alert_description, |
| 37 | "status": alert.alert.status, |
| 38 | "time_closed": alert.alert.time_closed, |
| 39 | "tags": [tag.tag.tag for tag in alert.alert.tags], |
| 40 | "assets": [ |
| 41 | { |
| 42 | "asset_name": asset.asset_name, |
| 43 | "agent_id": asset.agent_id, |
| 44 | } |
| 45 | for asset in alert.alert.assets |
| 46 | ], |
| 47 | "comments": [ |
| 48 | { |
| 49 | "comment": comment.comment, |
| 50 | "user_name": comment.user_name, |
| 51 | "created_at": comment.created_at, |
| 52 | } |
| 53 | for comment in alert.alert.comments |
| 54 | ], |
| 55 | "context": { |
| 56 | "source": alert.alert.assets[0].alert_context.source |
| 57 | if alert.alert.assets and alert.alert.assets[0].alert_context |
| 58 | else None, |
| 59 | "context": alert.alert.assets[0].alert_context.context |
| 60 | if alert.alert.assets and alert.alert.assets[0].alert_context |
| 61 | else None, |
| 62 | } |
| 63 | if alert.alert.assets |
| 64 | else None, |
| 65 | "iocs": [ |
| 66 | { |
| 67 | "ioc_value": ioc.ioc.value, |
| 68 | "ioc_type": ioc.ioc.type, |
| 69 | "ioc_description": ioc.ioc.description, |
| 70 | } |
| 71 | for ioc in alert.alert.iocs |
| 72 | ], |
| 73 | } |
| 74 | for alert in case.alerts |
| 75 | ], |
| 76 | }, |
| 77 | } |
| 78 | |
| 79 | |
| 80 | def render_html_template(template_path: str, context: Dict[str, Dict[str, str]]) -> str: |
| 81 | """Render the Jinja HTML template with the provided context.""" |
| 82 | template_dir = os.path.dirname(template_path) |
| 83 | template_name = os.path.basename(template_path) |
| 84 | |
| 85 | env = Environment(loader=FileSystemLoader(template_dir)) |
| 86 | template = env.get_template(template_name) |
| 87 | rendered_html = template.render(context) |
| 88 | |
| 89 | # Save rendered HTML to a temporary file |
| 90 | with NamedTemporaryFile(delete=False, suffix=".html") as tmp: |
| 91 | tmp.write(rendered_html.encode("utf-8")) |
| 92 | return tmp.name |
| 93 | |
| 94 | |
| 95 | def convert_html_to_pdf(html_path: str) -> str: |
| 96 | """Convert the HTML file to a PDF using wkhtmltopdf via pdfkit, with dynamic path detection for different platforms.""" |
| 97 | pdf_path = html_path.replace(".html", ".pdf") |
| 98 | wkhtmltopdf_paths = [] |
| 99 | |
| 100 | # Determine paths to wkhtmltopdf based on the current platform |
| 101 | try: |
| 102 | if platform.system() == "Windows": |
| 103 | # Common installation paths for wkhtmltopdf on Windows |
| 104 | wkhtmltopdf_paths = [ |
| 105 | r"C:\Program Files\wkhtmltopdf\bin\wkhtmltopdf.exe", |
| 106 | r"C:\Program Files (x86)\wkhtmltopdf\bin\wkhtmltopdf.exe", |
| 107 | ] |
| 108 | elif platform.system() == "Darwin": # macOS |
| 109 | # Common installation paths for wkhtmltopdf on macOS |
| 110 | wkhtmltopdf_paths = ["/usr/local/bin/wkhtmltopdf", "/opt/homebrew/bin/wkhtmltopdf"] # For macOS ARM (M1/M2) using Homebrew |
| 111 | elif platform.system() == "Linux": |
| 112 | # Common installation paths for wkhtmltopdf on Linux (Debian-based) |
| 113 | wkhtmltopdf_paths = ["/usr/bin/wkhtmltopdf", "/usr/local/bin/wkhtmltopdf"] |
| 114 | |
| 115 | # Try each path until a valid executable is found |
| 116 | path_to_wkhtmltopdf = None |
| 117 | for path in wkhtmltopdf_paths: |
| 118 | try: |
| 119 | # Check if the executable can be accessed |
| 120 | config = pdfkit.configuration(wkhtmltopdf=path) |
| 121 | path_to_wkhtmltopdf = path |
| 122 | break |
| 123 | except OSError: |
| 124 | continue |
| 125 | |
| 126 | # Raise an exception if no valid wkhtmltopdf path is found |
| 127 | if path_to_wkhtmltopdf is None: |
| 128 | raise FileNotFoundError("No valid wkhtmltopdf executable found. Ensure wkhtmltopdf is installed and accessible.") |
| 129 | |
| 130 | # Generate the PDF from HTML using the valid wkhtmltopdf path |
| 131 | pdfkit.from_file(html_path, pdf_path, configuration=config) |
| 132 | except Exception as e: |
| 133 | raise RuntimeError(f"Failed to convert HTML to PDF: {str(e)}") |
| 134 | |
| 135 | return pdf_path |
| 136 | |
| 137 | |
| 138 | def create_file_response_pdf(file_path: str, file_name: Optional[str] = "case_report.pdf") -> FileResponse: |
| 139 | """Create a FileResponse object for the rendered document.""" |
| 140 | return FileResponse( |
| 141 | file_path, |
| 142 | filename=file_name, |
| 143 | media_type="application/pdf", |
| 144 | ) |