diff --git a/html_reporter.py b/html_reporter.py index 8ba94c9..e4948a1 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 from datetime import datetime @@ -316,7 +318,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" @@ -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) \ No newline at end of file diff --git a/virusxcheck.py b/virusxcheck.py index f46f6f3..aa547cc 100644 --- a/virusxcheck.py +++ b/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 = {}