From 1bb20c6e2d1b026dfbec1856441e1038a7d5eb3f Mon Sep 17 00:00:00 2001 From: rawcell-agent Date: Sun, 12 Apr 2026 10:42:05 -0700 Subject: [PATCH] Add Python type hints throughout the codebase Adds type annotations to all public functions and methods across virusxcheck.py, html_reporter.py, and pdf_reporter.py for better IDE support, static analysis, and code clarity. Closes #5 Co-Authored-By: Claude Sonnet 4.6 --- html_reporter.py | 16 +++++++++------- pdf_reporter.py | 22 ++++++++++++---------- virusxcheck.py | 38 ++++++++++++++++++++------------------ 3 files changed, 41 insertions(+), 35 deletions(-) diff --git a/html_reporter.py b/html_reporter.py index 0d24c43..9a88fe0 100644 --- a/html_reporter.py +++ b/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 import pandas as pd @@ -309,7 +311,7 @@ """ -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" @@ -328,13 +330,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 = [] @@ -432,7 +434,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 = [] @@ -534,7 +536,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 = {} @@ -576,7 +578,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) @@ -623,7 +625,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) \ No newline at end of file diff --git a/pdf_reporter.py b/pdf_reporter.py index ff41f45..9e762e9 100644 --- a/pdf_reporter.py +++ b/pdf_reporter.py @@ -2,6 +2,8 @@ PDF Report Generator module for Virus.xcheck """ +from __future__ import annotations + import os from datetime import datetime from fpdf2 import FPDF @@ -14,13 +16,13 @@ from reportlab.lib.units import inch, mm -class VirusXcheckPDF(FPDF): - def __init__(self): +class VirusXcheckPDF(FPDF): + def __init__(self) -> None: super().__init__() self.WIDTH = 210 self.HEIGHT = 297 - - def header(self): + + def header(self) -> None: logo_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'logo.png') if os.path.exists(logo_path): self.image(logo_path, 10, 8, 33) @@ -36,7 +38,7 @@ def header(self): self.cell(0, 10, f'Generated on {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}', 0, 0, 'R') self.ln(20) - def footer(self): + def footer(self) -> None: """Create footer with page numbers""" self.set_y(-15) self.set_font('helvetica', 'I', 8) @@ -46,13 +48,13 @@ def footer(self): class PDFReporter: - def __init__(self): + def __init__(self) -> None: self.pdf = VirusXcheckPDF() self.pdf.set_auto_page_break(auto=True, margin=15) self.pdf.add_page() self.pdf.alias_nb_pages() - def add_summary_section(self, results): + def add_summary_section(self, results: dict) -> None: # Calculate summary statistics total_hashes = len(results) vx_found = sum(1 for details in results.values() if 'Found in VX database' in details['status']) @@ -132,7 +134,7 @@ def add_summary_section(self, results): self.pdf.ln(8) - def add_detailed_results(self, results): + def add_detailed_results(self, results: dict) -> None: # detailed results for each hash # Section header self.pdf.set_font('helvetica', 'B', 16) @@ -241,13 +243,13 @@ def add_detailed_results(self, results): self.pdf.line(10, self.pdf.get_y(), 200, self.pdf.get_y()) self.pdf.ln(8) - def generate_report(self, results, output_file): + def generate_report(self, results: dict, output_file: str) -> str: # Generate the PDF report self.add_summary_section(results) self.add_detailed_results(results) self.pdf.output(output_file) return output_file -def generate_pdf_report(results, output_file): +def generate_pdf_report(results: dict, output_file: str) -> str: reporter = PDFReporter() return reporter.generate_report(results, output_file) \ No newline at end of file diff --git a/virusxcheck.py b/virusxcheck.py index fe11a35..a6a5eeb 100644 --- a/virusxcheck.py +++ b/virusxcheck.py @@ -9,6 +9,8 @@ """ +from __future__ import annotations + import os import sys import warnings @@ -54,9 +56,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: @@ -75,14 +77,14 @@ def read_csv(file_path): 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: @@ -109,7 +111,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}" @@ -140,17 +142,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 @@ -167,7 +169,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 @@ -204,9 +206,9 @@ def extract_scan_results(self, vt_data): except Exception as e: return {"error": f"Error parsing VirusTotal data: {str(e)}"} -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 = {} # Validate hash length and type if len(hash_value) == 64: # SHA-256 @@ -250,7 +252,7 @@ def check_hash(hash_value, api, vt_api=None): else: return {"status": "Invalid hash length", "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) @@ -324,12 +326,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}") @@ -481,7 +483,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') @@ -512,7 +514,7 @@ def update_env_file(api_key): with open(env_file, 'w') as f: f.writelines(content) -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') @@ -544,7 +546,7 @@ def update_env_file_multiple(env_content): with open(env_file, 'w') as f: f.writelines(content) -def main(): +def main() -> None: # Restore stderr only for controlled output try: # First save original args for later use to avoid potential stderr output