scanner.py

import os
import hashlib
import psutil
import winreg
import requests
import logging
import yara
from pathlib import Path
import matplotlib.pyplot as plt

# Setup logging to save damaged area mapping
logging.basicConfig(
    filename="system_scan.log",
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s"
)

# VirusTotal API key (Replace with your API key)
VIRUSTOTAL_API_KEY = "your_api_key_here"
VIRUSTOTAL_URL = "https://www.virustotal.com/api/v3/files"

# YARA rules directory
YARA_RULES_DIR = "yara_rules"

# Visualization data
damaged_areas = []

def hash_file(file_path):
    """Calculate the SHA256 hash of a file."""
    try:
        with open(file_path, "rb") as f:
            file_hash = hashlib.sha256(f.read()).hexdigest()
        return file_hash
    except Exception as e:
        logging.error(f"Error hashing file {file_path}: {e}")
        return None

def scan_file_with_virustotal(file_hash):
    """Scan a file hash using the VirusTotal API."""
    try:
        headers = {"x-apikey": VIRUSTOTAL_API_KEY}
        response = requests.get(f"{VIRUSTOTAL_URL}/{file_hash}", headers=headers)
        if response.status_code == 200:
            data = response.json()
            positives = data.get("data", {}).get("attributes", {}).get("last_analysis_stats", {}).get("malicious", 0)
            return positives > 0
        else:
            logging.warning(f"VirusTotal API error: {response.status_code}")
            return False
    except Exception as e:
        logging.error(f"Error contacting VirusTotal: {e}")
        return False

def scan_with_yara(file_path):
    """Scan a file with YARA rules."""
    try:
        yara_rules = yara.compile(filepath=os.path.join(YARA_RULES_DIR, "rules.yara"))
        matches = yara_rules.match(file_path)
        if matches:
            logging.warning(f"YARA match found in file: {file_path}")
            return True
    except Exception as e:
        logging.error(f"Error running YARA on {file_path}: {e}")
    return False

def scan_files(directory):
    """Scan files in a directory."""
    logging.info(f"Scanning directory: {directory}")
    for root, _, files in os.walk(directory):
        for file in files:
            file_path = os.path.join(root, file)
            file_hash = hash_file(file_path)
            if file_hash:
                if scan_file_with_virustotal(file_hash) or scan_with_yara(file_path):
                    logging.warning(f"Malicious or suspicious file detected: {file_path}")
                    damaged_areas.append(file_path)

def scan_registry():
    """Scan Windows registry for suspicious entries."""
    try:
        suspicious_keys = []
        with winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, r"Software\Microsoft\Windows\CurrentVersion\Run") as key:
            i = 0
            while True:
                try:
                    value = winreg.EnumValue(key, i)
                    suspicious_keys.append(value)
                    i += 1
                except OSError:
                    break
        for key in suspicious_keys:
            logging.info(f"Suspicious registry entry: {key}")
            damaged_areas.append(f"Registry: {key}")
    except Exception as e:
        logging.error(f"Error scanning registry: {e}")

def scan_processes():
    """Scan running processes for anomalies."""
    logging.info("Scanning running processes.")
    for proc in psutil.process_iter(attrs=["pid", "name", "exe"]):
        try:
            process_name = proc.info["name"]
            process_exe = proc.info["exe"]
            if process_exe:
                file_hash = hash_file(process_exe)
                if file_hash and scan_file_with_virustotal(file_hash):
                    logging.warning(f"Malicious process detected: {process_name} ({process_exe})")
                    damaged_areas.append(f"Process: {process_name} ({process_exe})")
        except (psutil.AccessDenied, psutil.NoSuchProcess):
            continue

def visualize_damaged_areas():
    """Generate a visual map of damaged areas."""
    labels = [os.path.basename(area) for area in damaged_areas]
    values = [1] * len(damaged_areas)

    plt.figure(figsize=(10, 6))
    plt.barh(labels, values, color="red")
    plt.xlabel("Damage Count")
    plt.ylabel("Affected Items")
    plt.title("Visual Map of Damaged Areas")
    plt.tight_layout()
    plt.savefig("damaged_areas_map.png")
    plt.show()

def main():
    """Main function to run the scan."""
    # Create YARA rules directory if it doesn't exist
    os.makedirs(YARA_RULES_DIR, exist_ok=True)

    # Prompt for a directory to scan
    directory_to_scan = input("Enter the directory to scan (e.g., C:\\): ").strip()
    
    # Scan files
    scan_files(directory_to_scan)
    
    # Scan registry
    scan_registry()
    
    # Scan processes
    scan_processes()
    
    # Visualize damaged areas
    visualize_damaged_areas()
    
    logging.info("System scan completed. Check system_scan.log for details.")

if __name__ == "__main__":
    main()