feat: Refactored AI verdict module for better modularity and efficiency, removed unnecessary code, and improved overall structure.

Author: phantom0004

Main Files/modules/ai_verdict.py

@@ -0,0 +1,166 @@
+"""
+AI Verdict Processing Module
+Author: Phantom0004 (Daryl Gatt)
+
+Description:
+Handles AI-driven analysis of YARA scan results, providing structured and detailed insights into potential threats. 
+This module interacts with an external AI model to generate contextual intelligence based on detected Indicators of Compromise (IoCs). 
+
+Features:
+- **AI-Powered Verdicts:** Sends YARA scan results to an advanced AI model for context-aware threat analysis.
+- **Terminal Adaptive Formatting:** Dynamically adjusts output formatting based on terminal capabilities (plain text or Markdown).
+- **API Communication Handling:** Manages API requests and response handling, ensuring robust connectivity.
+
+Usage:
+- Pass structured YARA scan results to this module.
+- The AI model will process the structured IoC findings and return an actionable analysis.
+- Depending on terminal support, plain text and Markdown formatting will be selected.
+"""
+
+from rich.console import Console
+from rich.markdown import Markdown
+from typing import Tuple
+import requests
+
+class AIVerdict:
+    """
+    Processes YARA scan results and interacts with an AI model for threat analysis.
+
+    This class:
+    - Structures unstructured YARA scan results for AI-driven analysis.
+    - Categorizes YARA rules into general file analysis and malware-specific classifications.
+    - Formats AI-generated responses based on terminal support.
+    - Manages API requests and error handling for AI communication.
+
+    Attributes:
+        yara_results (list): The unstructured YARA scan results provided at initialization.
+        console (Console): Rich console instance for handling terminal output formatting.
+    """
+    
+    def __init__(self, structured_yara_results: list) -> dict:
+        self.structured_yara_results = structured_yara_results
+        self.console = Console()
+    
+    def generate_api_request(self) -> dict:  
+        """
+        Generates a structured API request for the AI model based on YARA scan results.
+
+        Args:
+            None
+
+        Returns:
+            dict: 
+                - A dictionary containing the structured payload for the AI model.
+                - The payload includes:
+                    - System and user messages with formatting based on terminal support.
+                    - YARA scan results formatted for analysis.
+                    - Model selection (`EleutherAI/gpt-j-6B`) for accurate responses.
+                    - A fixed seed value (`42`) to ensure deterministic output.
+        """
+        
+        # Define prompt for AI model - plaintext for basic terminals and markdown for advanced terminals
+        plaintext_prompt = """
+        Format responses as plain text with concise paragraphs and line breaks, ensuring maximum compatibility with all terminal environments. 
+        Avoid any Markdown, JSON, or structured formatting—just simple, unadorned text.
+        """
+        markdown_prompt = """
+        Format responses with rich Markdown formatting to enhance clarity, structure, and visual appeal. 
+        Incorporate headings, bullet lists, code blocks, and other Markdown elements to produce well-organized, visually engaging output suitable for terminal display.
+        """
+        
+        # Payload varies depending on terminal support for advanced formatting
+        payload = {
+            "messages": [
+                {
+                    "role": "system",
+                    "content": (
+                        "You are 'MORPHEUS_IQ', an advanced cybersecurity expert in malware analysis."
+                        "Provide precise and actionable insights from IoCs and signature findings, avoiding speculation, repetition, or hallucinations."
+                        "Keep responses short and terminal-friendly, summarizing only critical details for quick analysis."
+                        f"{markdown_prompt if self.supports_advanced_formatting() else plaintext_prompt}"
+                        "Ensure output is clear, well-structured, and formatted for immediate readability in a technical environment."
+                    )
+                },
+                {   
+                    "role": "user",
+                    "content": (
+                        "Analyze the provided scan results and deliver concise, actionable insights based on matched YARA signatures."
+                        f"\nScan Output: {str(self.structured_yara_results)}"
+                    )
+                }
+            ],
+            # Utalizes a high paramater model for more accurate results
+            "model": "EleutherAI/gpt-j-6B",
+            # This seed value ensures that the results are mostly deterministic
+            "seed": 42
+        }
+                
+        return payload
+            
+    def supports_advanced_formatting(self) -> bool:
+        """
+        Checks if the terminal supports advanced formatting for markdown rendering.
+
+        Args:
+            None
+
+        Returns:
+            bool:
+                - True if the terminal supports advanced formatting.
+                - False otherwise.
+                - This is determined by checking if the console is a terminal 
+                and if a color system is available.
+        """
+        
+        return self.console.is_terminal and self.console.color_system is not None
+
+    def format_string_to_markdown(self, ai_verdict_output) -> None:
+        """
+        Renders the AI-generated verdict as Markdown in the terminal.
+
+        Args:
+            ai_verdict_output (str): The AI-generated analysis output to be formatted and displayed.
+
+        Returns:
+            None:
+                - Outputs the formatted AI response directly to the console.
+                - Uses the `Markdown` class to render the response for enhanced readability.
+                - Requires a terminal that supports rich text formatting.
+                - If the terminal does not support advanced formatting, consider using plain text output.
+        """
+        
+        self.console.print(Markdown(ai_verdict_output.strip()))
+       
+    @staticmethod
+    def send_api_request(payload: dict) -> Tuple[str, str]:
+        """
+        Sends a POST request to the Pollinations API with the provided payload.
+
+        Args:
+            payload (dict): The JSON-formatted data to be sent in the API request.
+
+        Returns:
+            Tuple[str, str]:
+                - The API response text and a status indicator.
+                - If successful (`200 OK`), returns the API response text and `"success"`.
+                - If the request times out, returns an error message and `"fail"`.
+                - If a connection failure occurs, returns an error message detailing the issue and `"fail"`.
+                - If an unknown exception occurs, returns a generic error message and `"fail"`.
+                - If the API is offline (non-200 response), returns an appropriate message and `"fail"`.
+        """
+        
+        # Send POST request to Pollinations API
+        try:
+            response = requests.post("https://text.pollinations.ai/", headers={"Content-Type": "application/json"}, json=payload)
+        except requests.exceptions.Timeout:
+            return "[-] API Endpoint Timed Out! Please try again at a later time.", "fail"
+        except requests.exceptions.RequestException as error:
+            return f"[-] Critical Connection Failure. Unable to connect to API - Error : {error}", "fail"
+        except Exception as error:
+            return f"[-] An Unknown Error has Occured - Error : {error}", "fail"
+        
+        # Handle API response
+        if response.status_code == 200:
+            return response.text, "success"
+        else:
+            return f"[-] API Endpoint seems to be offline ... Please try again at a later time.", "fail"

Main Files/modules/AnalysisReportPDF.py Main Files/modules/analysis_report.pyMain Files/modules/AnalysisReportPDF.py renamed to Main Files/modules/analysis_report.py

@@ -21,16 +21,15 @@
 
 try:
     import vt
+    import requests
     from termcolor import colored
 except ModuleNotFoundError:
     exit(f"'virus_total.py' Library Error -> Please install all requirements via the 'requirements.txt' file with PIP")
 
-import hashlib
-import requests 
+from urllib.parse import urlparse, quote
+import hashlib 
 import json
 import os
-import urllib.parse
-from urllib.parse import urlparse
 
 class VirusTotalAPI:
     def __init__ (self, choice="1", data="placeholder", API_KEY="", client_obj=""):
@@ -368,7 +367,7 @@ def extract_behaviour_techniques_of_file(self):
 
     # Rescan URL if any fail arises
     def rescan_url(self):
-        data = urllib.parse.quote(self.data, safe='')
+        data = quote(self.data, safe='')
 
         # May not always work, but if it does, a re-scan will be submitted
         self.send_api_request_using_requests(f"urls/{data}/analyse")

Main Files/modules/virus_total.py

@@ -16,18 +16,11 @@
 """
 
 import os
+import yara
 import concurrent.futures
 import threading
 from typing import Union, List
-try:
-    import yara
-    from termcolor import colored
-except Exception as err:
-    if "libyara" in str(err):
-        print("\nLibyara not found in your 'Yara' installation. Please try uninstall all python dependencies and re-install them.")
-        exit("\nIf all else fails, open an issue on Morpheus with your details.")
-    else:
-        exit(f"'yara_analysis.py' Library Error -> Error on Import : {err}")
+from termcolor import colored
 
 # Parent class which handles the core of this file
 class BaseDetection:
@@ -124,17 +117,12 @@ def identify_matches(self, yara_object:yara.Rules) -> Union[List[yara.Match], st
     def output_yara_matches(yara_match:yara.Match) -> None:
         if yara_match:
             for match in yara_match:
-                # Required Variables
                 tags = f"Matched Tags: {str(match.tags)}" if match.tags else ""
                 metadata = f"Rule Description: {str(match.meta['description'])}" if match.meta["description"] else ""
                 print(' ' * 80, end='\r')  # Clears reminants from the dynamic printing            
-                
-                if "KRYPTOS" in str(match): 
-                    # Detect for KRYPTOS ransomware  
-                    print(f'[+] KRYPTOS RANSOMWARE DETECTED: {colored(str(match), attrs=["bold"])}')
-                else:
-                    # Detect for any other malware
-                    print(f"[+] Matched Rule: '{colored(str(match), 'green', attrs=['bold'])}'"+"  "+tags)
+            
+                # Detect for any other malware
+                print(f"[+] Matched Rule: '{colored(str(match), 'green', attrs=['bold'])}'"+"  "+tags)
 
                 # Metadata Matches
                 print(colored(metadata.capitalize(), "yellow"))