Add Python type hints throughout the codebase

html_reporter.py

@@ -3,6 +3,8 @@
 This module generates interactive HTML reports with charts for hash analysis results
 """
 
+from __future__ import annotations
+
 import os
 import json
 from datetime import datetime
@@ -316,7 +318,7 @@
 </html>"""
 
 
-def jinja_filesizeformat_filter(value, binary=False):
+def jinja_filesizeformat_filter(value: float | str | None, binary: bool = False) -> str:
     """Format file sizes for Jinja2 template"""
     if value is None or value == 'N/A':
         return "Unknown"
@@ -335,13 +337,13 @@ def jinja_filesizeformat_filter(value, binary=False):
 
 
 class HTMLReporter:
-    def __init__(self):
+    def __init__(self) -> None:
         """Initialize the HTML reporter"""
         self.env = Environment(autoescape=True)
         self.env.filters['filesizeformat'] = jinja_filesizeformat_filter
         self.template = self.env.from_string(HTML_TEMPLATE)
 
-    def create_detection_chart(self, results):
+    def create_detection_chart(self, results: dict) -> str:
         """Create a bar chart for detection rates"""
         data = []
 
@@ -439,7 +441,7 @@ def create_detection_chart(self, results):
 
         return chart_js
 
-    def create_distribution_chart(self, results):
+    def create_distribution_chart(self, results: dict) -> str:
         """Create a gauge chart for detection rates distribution, providing better visualization of malicious presence"""
         detection_rates = []
 
@@ -541,7 +543,7 @@ def create_distribution_chart(self, results):
 
         return chart_js
 
-    def create_tags_chart(self, results):
+    def create_tags_chart(self, results: dict) -> str:
         """Create a bar chart for common tags"""
         # Collect all tags
         all_tags = {}
@@ -583,7 +585,7 @@ def create_tags_chart(self, results):
 
         return chart_js
 
-    def generate_report(self, results, output_file):
+    def generate_report(self, results: dict, output_file: str) -> str:
         """Generate HTML report with interactive charts"""
         # Generate charts
         detection_chart_js = self.create_detection_chart(results)
@@ -630,7 +632,7 @@ def generate_report(self, results, output_file):
 
 
 # Function to use from main application
-def generate_html_report(results, output_file):
+def generate_html_report(results: dict, output_file: str) -> str:
     """Generate an HTML report with interactive charts from results"""
     reporter = HTMLReporter()
     return reporter.generate_report(results, output_file)

virusxcheck.py

@@ -9,6 +9,8 @@
 
 """
 
+from __future__ import annotations
+
 __version__ = "0.2.1"
 
 import os
@@ -57,9 +59,9 @@
 DEFAULT_API_KEY = os.getenv("VIRUSXCHECK_API_KEY", "")
 DEFAULT_VT_API_KEY = os.getenv("VIRUSTOTAL_API_KEY", "")
 
-def read_csv(file_path):
+def read_csv(file_path: str) -> list[str]:
     # Extract SHA256 hashes from a CSV file using regex
-    hashes = []
+    hashes: list[str] = []
     hex_pattern = re.compile(r'\b[a-fA-F0-9]{64}\b')  # Pattern to match SHA256
     try:
         with open(file_path, mode='r', newline='', encoding='utf-8') as file:
@@ -78,14 +80,14 @@ def read_csv(file_path):
         sys.exit(1)
 
 class VirusExchangeAPI:
-    def __init__(self, api_key):
+    def __init__(self, api_key: str) -> None:
         self.api_key = api_key
         self.session = requests.Session()
         self.session.headers.update({"Authorization": f"Bearer {api_key}"})
-    
+
     @sleep_and_retry
     @limits(calls=15, period=1)
-    def get_sample_details(self, sha256_hash):
+    def get_sample_details(self, sha256_hash: str) -> dict:
         # Get sample details from the Virus.Exchange API
         url = f"{API_BASE_URL}/samples/{sha256_hash}"
         try:
@@ -112,7 +114,7 @@ def get_sample_details(self, sha256_hash):
                 "virustotal_url": f"https://www.virustotal.com/gui/file/{sha256_hash}"
             }
 
-    def fallback_check(self, sha256_hash):
+    def fallback_check(self, sha256_hash: str) -> dict:
         # Fallback method to check S3 bucket directly if the API fails
         vx_url = f"https://s3.us-east-1.wasabisys.com/vxugmwdb/{sha256_hash}"
         virustotal_url = f"https://www.virustotal.com/gui/file/{sha256_hash}"
@@ -143,17 +145,17 @@ def fallback_check(self, sha256_hash):
             }
 
 class VirusTotalAPI:
-    def __init__(self, api_key):
+    def __init__(self, api_key: str) -> None:
         self.api_key = api_key
         self.session = requests.Session()
         self.session.headers.update({
             "x-apikey": api_key,
             "Accept": "application/json"
         })
-    
+
     @sleep_and_retry
     @limits(calls=4, period=60)  # VT API rate limits: 4 requests per minute for standard API keys
-    def get_file_report(self, file_hash):
+    def get_file_report(self, file_hash: str) -> dict | None:
         # Get detailed file information from the VirusTotal API
         if not self.api_key:
             return None
@@ -170,7 +172,7 @@ def get_file_report(self, file_hash):
         except requests.RequestException as e:
             return {"error": f"Request Error: {e}"}
 
-    def extract_scan_results(self, vt_data):
+    def extract_scan_results(self, vt_data: dict | None) -> dict | None:
         # Pull out the important bits from the VirusTotal API response
         if not vt_data or "error" in vt_data:
             return vt_data
@@ -207,7 +209,7 @@ def extract_scan_results(self, vt_data):
         except Exception as e:
             return {"error": f"Error parsing VirusTotal data: {str(e)}"}
 
-def validate_hash(hash_value):
+def validate_hash(hash_value: str) -> str | None:
     """Validate that a hash is strictly hex characters of a known length.
     Returns the hash type string or None if invalid."""
     if not re.fullmatch(r'[a-fA-F0-9]+', hash_value):
@@ -223,9 +225,9 @@ def validate_hash(hash_value):
         return "sha512"
     return None
 
-def check_hash(hash_value, api, vt_api=None):
+def check_hash(hash_value: str, api: VirusExchangeAPI, vt_api: VirusTotalAPI | None = None) -> dict:
     # Checks a hash using the Virus.Exchange API with fallback and VirusTotal lookup if available
-    result = {}
+    result: dict = {}
 
     hash_type = validate_hash(hash_value)
     if hash_type is None:
@@ -273,7 +275,7 @@ def check_hash(hash_value, api, vt_api=None):
     else:
         return {"status": "Invalid hash format", "virustotal_url": None}
 
-def write_to_csv(file_path, data):
+def write_to_csv(file_path: str, data: dict) -> None:
     # Export results to a CSV file with all the metadata we've collected
     with open(file_path, 'w', newline='', encoding='utf-8') as file:
         writer = csv.writer(file)
@@ -347,12 +349,12 @@ def write_to_csv(file_path, data):
                 vt_tags
             ])
 
-def write_to_json(file_path, data):
+def write_to_json(file_path: str, data: dict) -> None:
     # Save all the collected data in a nicely formatted JSON file
     with open(file_path, 'w', encoding='utf-8') as f:
         json.dump(data, f, indent=4)
 
-def pretty_print_results(results):
+def pretty_print_results(results: dict) -> None:
     # Format and display the results with nice colors and clear organization
     print()
     print(f"{Fore.CYAN}{'═'*80}{Style.RESET_ALL}")
@@ -504,7 +506,7 @@ def pretty_print_results(results):
 
         print(f"{Fore.CYAN}{'─'*80}{Style.RESET_ALL}\n")
 
-def update_env_file(api_key):
+def update_env_file(api_key: str) -> None:
     # Saves a new API key to the .env file for future use
     env_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), '.env')
 
@@ -541,7 +543,7 @@ def update_env_file(api_key):
     except OSError:
         pass
 
-def update_env_file_multiple(env_content):
+def update_env_file_multiple(env_content: dict[str, str]) -> None:
     # Batch updates to the .env file with multiple API keys
     env_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), '.env')
 
@@ -582,7 +584,7 @@ def update_env_file_multiple(env_content):
     except OSError:
         pass
 
-def main():
+def main() -> None:
     parser = argparse.ArgumentParser(description='Virus.xcheck CLI Tool')
     parser.add_argument('-f', '--file', help='Path to CSV file containing hashes')
     parser.add_argument('-o', '--output', help='Path to output file (CSV or JSON format)')
@@ -591,19 +593,19 @@ def main():
                         help='Interactively save API keys to .env file')
     parser.add_argument('--no-color', action='store_true', help='Disable colored output')
     parser.add_argument('--html', help='Generate HTML report with interactive charts')
-    
+
     args = parser.parse_args()
-    
+
     # Check if no arguments were provided
     if len(sys.argv) == 1:
         print(__doc__)  # Print ASCII art
         parser.print_help()
         sys.exit(1)
-    
+
     # Disable colors if requested
     if args.no_color:
         colorama.deinit()
-    
+
     # Save API keys interactively if requested
     if args.save_config:
         env_content = {}