main.py

import wmi
import hashlib
import requests
import os
import json
import time
import sys

# Configuration file for storing the API key
CONFIG_FILE = "config.json"

def load_config():
    """Loads configuration from config.json."""
    if os.path.exists(CONFIG_FILE):
        try:
            with open(CONFIG_FILE, "r") as f:
                return json.load(f)
        except Exception as e:
            print(f"Error loading configuration: {e}")
            sys.exit(1)
    else:
        print(f"Configuration file {CONFIG_FILE} not found!")
        sys.exit(1)

config = load_config()
API_KEY = config.get("API_KEY")
if not API_KEY:
    print("No API_KEY found in the configuration!")
    sys.exit(1)

# Additional configuration settings
THRESHOLD = 1                       # At least 1 malicious report qualifies as a threat
CACHE_FILE = "vt_cache.json"
SLEEP_INTERVAL = 15                 # Delay (in seconds) between API queries

def load_cache():
    """Loads the cache from the JSON file, if it exists."""
    if os.path.exists(CACHE_FILE):
        try:
            with open(CACHE_FILE, "r") as f:
                return json.load(f)
        except Exception as e:
            print(f"Error loading cache: {e}")
    return {}

def save_cache(cache):
    """Saves the cache to the JSON file."""
    try:
        with open(CACHE_FILE, "w") as f:
            json.dump(cache, f, indent=4)
    except Exception as e:
        print(f"Error saving cache: {e}")

def get_processes_with_exe():
    """
    Retrieves all running processes and returns a list of dictionaries,
    each containing the PID, process name, and path to the executable.
    """
    c = wmi.WMI()
    processes = []
    for process in c.Win32_Process():
        if process.ExecutablePath and os.path.exists(process.ExecutablePath):
            processes.append({
                "pid": process.ProcessId,
                "name": process.Name,
                "exe": process.ExecutablePath
            })
    return processes

def calculate_file_hash(file_path, algorithm='sha256'):
    """Calculates the hash (e.g., SHA-256) of the specified file."""
    hash_func = hashlib.new(algorithm)
    try:
        with open(file_path, 'rb') as f:
            for chunk in iter(lambda: f.read(4096), b''):
                hash_func.update(chunk)
        return hash_func.hexdigest()
    except Exception as e:
        print(f"Error reading file {file_path}: {e}")
        return None

def check_file_with_virustotal(file_hash, api_key):
    """
    Queries the VirusTotal API using the file hash and returns the number
    of malicious reports and the total number of reports.
    If a 401 error occurs, the script terminates.
    """
    url = "https://www.virustotal.com/api/v3/files/" + file_hash
    headers = {"x-apikey": api_key}
    try:
        response = requests.get(url, headers=headers)
    except Exception as e:
        print(f"Error querying VirusTotal: {e}")
        return None, None

    if response.status_code == 200:
        data = response.json()
        stats = data["data"]["attributes"]["last_analysis_stats"]
        malicious = stats.get("malicious", 0)
        total = sum(stats.values())
        return malicious, total
    elif response.status_code == 401:
        print("VirusTotal API Error: Status Code 401. Please check your API key!")
        sys.exit(1)
    else:
        print(f"VirusTotal API Error: Status Code {response.status_code}")
        return None, None

def format_seconds(seconds):
    """Formats seconds as hh:mm:ss."""
    return time.strftime("%H:%M:%S", time.gmtime(seconds))

def scan_processes_for_threats(api_key, cache, threshold=THRESHOLD):
    """
    Retrieves all processes, computes their file hashes, and builds a list.
    For each process whose hash is not yet in the cache, the VirusTotal API is queried
    (with a SLEEP_INTERVAL pause per API call). The progress display shows:
      - Number of API calls completed (out of the total required)
      - Estimated remaining time (based on remaining API calls)
      - Number of processes classified as threats versus clean
    Finally, a detailed list of the processes flagged as malicious is printed.
    """
    processes = get_processes_with_exe()
    process_list = []

    # Pre-calculate the hash for each process and add to the list
    for proc in processes:
        file_hash = calculate_file_hash(proc["exe"])
        if file_hash:
            proc["file_hash"] = file_hash
            process_list.append(proc)
    
    total_processes = len(process_list)
    if total_processes == 0:
        print("No processes with executables found.")
        return

    # Determine how many API queries are needed (i.e., for processes whose hash is not cached)
    total_api_calls = sum(1 for proc in process_list if proc["file_hash"] not in cache)
    api_remaining = total_api_calls
    api_calls_done = 0

    print(f"{total_api_calls} API queries will be performed (one every {SLEEP_INTERVAL} seconds).")

    threat_count = 0
    clean_count = 0
    malicious_processes = []

    for proc in process_list:
        file_hash = proc["file_hash"]
        # Use cached data if available
        if file_hash in cache:
            result = cache[file_hash]
            mal_votes = result.get("malicious", 0)
            tot_votes = result.get("total", 0)
        else:
            mal_votes, tot_votes = check_file_with_virustotal(file_hash, api_key)
            if mal_votes is not None:
                cache[file_hash] = {"malicious": mal_votes, "total": tot_votes}
                # Immediately save the cache so progress is retained
                save_cache(cache)
            time.sleep(SLEEP_INTERVAL)
            api_calls_done += 1
            api_remaining = total_api_calls - api_calls_done

        if mal_votes is None:
            # Skip this process if a valid response wasn't received
            continue

        if mal_votes >= threshold:
            threat_count += 1
            malicious_processes.append(proc)
        else:
            clean_count += 1

        # Calculate estimated remaining time based on remaining API calls
        estimated_seconds = api_remaining * SLEEP_INTERVAL
        remaining_str = format_seconds(estimated_seconds)

        # Update progress display (single-line update)
        progress_line = (f"[{api_calls_done}/{total_api_calls}] Estimated time: {remaining_str} | "
                         f"Threats found: {threat_count} | Clean: {clean_count}")
        sys.stdout.write("\r" + progress_line)
        sys.stdout.flush()

    print("\n\nScan complete.")
    print(f"Total processes scanned: {total_processes}")
    print(f"Threats found: {threat_count}")
    print(f"Clean processes: {clean_count}")
    if threat_count > 0:
        print("\nList of potentially malicious processes:")
        for proc in malicious_processes:
            print(f"PID {proc['pid']}: {proc['name']} - {proc['exe']}")

def main():
    print("Starting process scanner PyPScan...\n")
    cache = load_cache()
    scan_processes_for_threats(API_KEY, cache)
    save_cache(cache)
    print("\nCache updated.")

if __name__ == "__main__":
    main()