minor update

src/engine_b/model.py

@@ -1,42 +1,46 @@
+"""
+Contains the CNN model for malware classification.
+"""
+
 import torch.nn as nn
 
+
 class MalwareCNN(nn.Module):
     """
     Convolutional Neural Network for classifying malware families based on byte images.
-    
+
     Args:
         num_classes (int): Number of unique malware families to classify.
     """
+
     def __init__(self, num_classes=24):
         super(MalwareCNN, self).__init__()
         self.features = nn.Sequential(
             nn.Conv2d(1, 32, kernel_size=3, padding=1),
             nn.ReLU(),
-            nn.MaxPool2d(2, 2), # 128 -> 64
-            
+            nn.MaxPool2d(2, 2),  # 128 -> 64
             nn.Conv2d(32, 64, kernel_size=3, padding=1),
             nn.ReLU(),
-            nn.MaxPool2d(2, 2), # 64 -> 32
-            
+            nn.MaxPool2d(2, 2),  # 64 -> 32
             nn.Conv2d(64, 128, kernel_size=3, padding=1),
             nn.ReLU(),
-            nn.MaxPool2d(2, 2), # 32 -> 16
+            nn.MaxPool2d(2, 2),  # 32 -> 16
         )
         self.classifier = nn.Sequential(
             nn.Flatten(),
             nn.Linear(128 * 16 * 16, 512),
             nn.ReLU(),
-            nn.Dropout(0.6), # Increased dropout to prevent overfitting
-            nn.Linear(512, num_classes)
+            nn.Dropout(0.6),  # Increased dropout to prevent overfitting
+            nn.Linear(512, num_classes),
         )
 
     def forward(self, x):
         """
         Forward pass through the CNN.
-        
+
         Args:
             x (torch.Tensor): Input batch of images.
-        
+
         Returns:
             torch.Tensor: Logits for each class.
         """

src/ui/app.py

@@ -62,7 +62,10 @@ def analyze_malware(file_path):
 
         # Ensemble Logic
         ensemble_score = (prob_a + result_b["confidence"] * 100) / 2
-        ensemble_verdict = f"Overall Risk Score: {ensemble_score:.1f}/100"
+        final_verdict_text = "MALICIOUS" if ensemble_score > 50 else "BENIGN"
+        ensemble_verdict = f"Final Verdict: {final_verdict_text} (Risk Score: {ensemble_score:.1f}/100)"
+
+        top_3_probs = sorted(all_probs.items(), key=lambda x: x[1], reverse=True)[:3]
 
         # Generate PDF Report
         pdf_path = generate_pdf_report(
@@ -74,6 +77,8 @@ def analyze_malware(file_path):
             prob_b=result_b["confidence"] * 100,
             image_b=image_b,
             ensemble_score=ensemble_score,
+            final_verdict=final_verdict_text,
+            top_3_probs=top_3_probs,
         )
 
         return (
@@ -97,7 +102,16 @@ def analyze_malware(file_path):
 
 
 def generate_pdf_report(
-    file_path, status_a, prob_a, top_features, family_b, prob_b, image_b, ensemble_score
+    file_path,
+    status_a,
+    prob_a,
+    top_features,
+    family_b,
+    prob_b,
+    image_b,
+    ensemble_score,
+    final_verdict,
+    top_3_probs,
 ):
 
     class PDFReport(FPDF):
@@ -111,12 +125,6 @@ def generate_pdf_report(
             self.cell(0, 10, "MALWARE ANALYSIS REPORT", border=0, align="C")
             self.ln(15)
 
-        def footer(self):
-            self.set_y(-15)
-            self.set_font("Helvetica", style="I", size=8)
-            self.set_text_color(128, 128, 128)
-            self.cell(0, 10, f"Page {self.page_no()}", align="C")
-
     pdf = PDFReport()
     pdf.add_page()
     pdf.set_auto_page_break(auto=True, margin=15)
@@ -157,11 +165,16 @@ def generate_pdf_report(
     )
 
     pdf.set_font("Helvetica", size=11)
-    pdf.cell(40, 8, "Risk Score:", border=0)
+    pdf.cell(40, 8, "Final Verdict:", border=0)
     pdf.set_font("Helvetica", style="B", size=12)
-    # Highlight risk score in text
+
     pdf.cell(
-        0, 8, f"{ensemble_score:.1f} / 100", border=0, new_x="LMARGIN", new_y="NEXT"
+        0,
+        8,
+        f"{final_verdict} (Risk Score: {ensemble_score:.1f} / 100)",
+        border=0,
+        new_x="LMARGIN",
+        new_y="NEXT",
     )
     pdf.ln(5)
 
@@ -253,6 +266,39 @@ def generate_pdf_report(
     pdf.rect(x_pos, y_pos, 70, 70)
     pdf.image(img_path, x=x_pos, y=y_pos, w=70, h=70)
 
+    # Draw Top 3 Probabilities Table next to the image
+    pdf.set_y(y_pos)
+
+    # Table Header
+    pdf.set_x(x_pos + 75)
+    pdf.set_font("Helvetica", style="B", size=10)
+    pdf.cell(70, 8, "Predicted Family", border=1, fill=True)
+    pdf.cell(
+        40,
+        8,
+        "Confidence",
+        border=1,
+        new_x="LMARGIN",
+        new_y="NEXT",
+        fill=True,
+        align="C",
+    )
+
+    # Table Rows
+    pdf.set_font("Helvetica", size=10)
+    for family, prob in top_3_probs:
+        pdf.set_x(x_pos + 75)
+        pdf.cell(70, 8, f" {family}", border=1)
+        pdf.cell(
+            40,
+            8,
+            f"{prob * 100:.2f}%",
+            border=1,
+            new_x="LMARGIN",
+            new_y="NEXT",
+            align="C",
+        )
+
     report_path = os.path.join(
         temp_dir, f"Analysis_Report_{os.path.basename(file_path)}.pdf"
     )