Update app.py

app.py

@@ -1,8 +1,13 @@
 import gradio as gr
 from analyze_email_main import analyze
-import pdfkit
+from weasyprint import HTML
 import tempfile
-import os
+
+def html_to_pdf(html_content):
+    """Convert HTML string to a temporary PDF file and return its path."""
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_pdf:
+        HTML(string=html_content).write_pdf(tmp_pdf.name)
+        return tmp_pdf.name
 
 def analyze_email(file_obj, export_pdf=False):
     if file_obj is None:
@@ -11,15 +16,15 @@ def analyze_email(file_obj, export_pdf=False):
     results = analyze(file_obj.name)
 
     # Basic verdict info
-    attack_score = results[0]  # e.g., "Attack Score: 85"
-    attack_type = results[1]   # e.g., "Attack Type: Credential Harvesting"
-    verdict = results[2]       # e.g., "Final Verdict: Malicious"
+    attack_score = results[0]
+    attack_type = results[1]
+    verdict = results[2]
 
-    score_val = int("".join(filter(str.isdigit, attack_score)))  # extract numeric score
+    score_val = int("".join(filter(str.isdigit, attack_score)))
     if score_val > 100:  
         score_val = 100
 
-    # Assign verdict color
+    # Verdict color
     if "Malicious" in verdict:
         verdict_color = "#e74c3c"  # red
     elif "Suspicious" in verdict:
@@ -29,7 +34,7 @@ def analyze_email(file_obj, export_pdf=False):
     else:
         verdict_color = "#2ecc71"  # green
 
-    # Extract report sections
+    # Extract sections
     try:
         tags_index = results.index("---- Attack Analysis Tags ----")
         findings_index = results.index("---- Detailed Findings ----")
@@ -42,7 +47,7 @@ def analyze_email(file_obj, export_pdf=False):
     findings = results[findings_index + 1 : body_index]
     highlighted_body = results[body_index + 1 :]
 
-    # Styled HTML Report
+    # HTML Report
     html = f"""
     <div style='font-family: Arial, sans-serif; padding: 20px; max-width: 800px; margin:auto;'>
 
@@ -80,10 +85,7 @@ def analyze_email(file_obj, export_pdf=False):
 
     pdf_path = None
     if export_pdf:
-        # Generate temporary PDF file
-        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_pdf:
-            pdf_path = tmp_pdf.name
-        pdfkit.from_string(html, pdf_path)
+        pdf_path = html_to_pdf(html)
 
     return html, pdf_path