Update app.py

app.py

@@ -20,6 +20,8 @@ from transformers import (
     AutoModelForSequenceClassification,
     T5ForConditionalGeneration,
 )
+from reportlab.platypus import SimpleDocTemplate, Paragraph
+from reportlab.lib.styles import getSampleStyleSheet
 from huggingface_hub import hf_hub_download
 import numpy as np
 
@@ -353,7 +355,23 @@ def get_json_report(code):
     if not code or not code.strip(): return {"error": "No code provided"}
     return build_json_report(code)
 
+def create_pdf(report_text):
+    pdf_path = "security_report.pdf"
+
+    doc = SimpleDocTemplate(pdf_path)
+
+    styles = getSampleStyleSheet()
+
+    elements = [
+        Paragraph(
+            report_text.replace("\n", "<br/>"),
+            styles["BodyText"]
+        )
+    ]
+
+    doc.build(elements)
 
+    return pdf_path
 EXAMPLES = [
     ["""import sqlite3\n\ndef get_user(username):\n    conn = sqlite3.connect('users.db')\n    query = f"SELECT * FROM users WHERE username = '{username}'"\n    return conn.execute(query).fetchone()\n"""],
     ["""#include <stdio.h>\n#include <string.h>\n\nvoid process_input(char *user_input) {\n    char buffer[64];\n    strcpy(buffer, user_input);\n    printf("Processed: %s\\n", buffer);\n}\n"""],
@@ -385,16 +403,37 @@ with gr.Blocks(
                 analyze_btn = gr.Button("🔍 Analyze Security", variant="primary", size="lg")
                 json_btn = gr.Button("📋 JSON Report", variant="secondary", size="lg")
         with gr.Column(scale=1):
-            report_output = gr.Markdown(label="Security Report")
-            json_output = gr.JSON(label="JSON Report", visible=False)
+    report_output = gr.Markdown(label="Security Report")
+
+    download_btn = gr.Button(
+        "📄 Generate PDF Report"
+    )
+
+    pdf_file = gr.File(
+        label="Download Report"
+    )
+
+    json_output = gr.JSON(
+        label="JSON Report",
+        visible=False
+    )
 
     gr.Examples(examples=EXAMPLES, inputs=[code_input], label="Example Code Snippets")
 
     def show_json(code):
         return gr.update(visible=True, value=get_json_report(code))
 
-    analyze_btn.click(fn=analyze_code, inputs=[code_input], outputs=[report_output], api_name="analyze")
-    json_btn.click(fn=show_json, inputs=[code_input], outputs=[json_output])
+    analyze_btn.click(
+        fn=analyze_code, inputs=[code_input], outputs=[report_output], api_name="analyze"
+                     )
+    json_btn.click(
+        fn=show_json, inputs=[code_input], outputs=[json_output]
+    )
+    download_btn.click(
+    fn=create_pdf,
+    inputs=[report_output],
+    outputs=[pdf_file]
+)
 
     with gr.Row(visible=False):
         api_json_btn = gr.Button("get_json", visible=False)
@@ -431,91 +470,7 @@ curl -X POST https://ayshajavd-code-security-analyzer.hf.space/call/analyze \\
 | **A08: Integrity Failures** | CWE-502 |
 | **A10: SSRF** | CWE-918 |
     """)
-def generate_report(
-    code,
-    prediction,
-    cwe,
-    owasp,
-    severity,
-    explanation,
-    recommendations,
-    fixed_code
-):
-    report = f"""
-SOFTWARE SECURITY ASSESSMENT REPORT
-
-Analysis Result
----------------
-Prediction: {prediction}
-
-Detected Vulnerability
-----------------------
-CWE: {cwe}
-OWASP: {owasp}
-Severity: {severity}
-
-Explanation
------------
-{explanation}
-
-Recommendations
----------------
-{recommendations}
-
-Suggested Secure Version
-------------------------
-{fixed_code}
-
-Original Code
--------------
-{code}
-"""
 
-    return report
-    report_output = gr.Textbox(
-    label="Security Report",
-    lines=25
-)
-report = generate_report(
-    code,
-    prediction,
-    cwe,
-    owasp,
-    severity,
-    explanation,
-    recommendations,
-    fixed_code
-)
-return (
-    prediction,
-    explanation,
-    fixed_code,
-    report
-)
-from reportlab.platypus import SimpleDocTemplate, Paragraph
-from reportlab.lib.styles import getSampleStyleSheet
-
-def create_pdf(report_text):
-    pdf_path = "security_report.pdf"
-
-    doc = SimpleDocTemplate(pdf_path)
-
-    styles = getSampleStyleSheet()
-
-    elements = [
-        Paragraph(report_text.replace("\n", "<br/>"), styles["BodyText"])
-    ]
-
-    doc.build(elements)
-
-    return pdf_path
-download_btn = gr.Button("Generate PDF Report")
-pdf_file = gr.File(label="Download Report")
-download_btn.click(
-    create_pdf,
-    inputs=report_output,
-    outputs=pdf_file
-)
 
 if __name__ == "__main__":
     demo.launch()