malware-scanner.py

import os
import sys
import json
import requests
import argparse
import signal # NEW: To catch the interrupt signal
from kubernetes import client, config
from kubernetes.stream import stream
from collections import Counter

# ANSI Escape Sequences
RED = "\033[91m"
YELLOW = "\033[93m"
GREEN = "\033[92m"
BOLD = "\033[1m"
RESET = "\033[0m"

# Global flag for graceful shutdown
interrupted = False

def signal_handler(sig, frame):
    """Callback for Ctrl+C to set the interrupt flag."""
    global interrupted
    if not interrupted:
        print(f"\n\n{YELLOW}[!] Interrupt received. Stopping scan and finalizing report...{RESET}")
        interrupted = True

# Register the signal handler
signal.signal(signal.SIGINT, signal_handler)

# --- API Interaction Logic ---

def get_osv_details(vuln_id):
    if interrupted: return None
    url = f"https://api.osv.dev/v1/vulns/{vuln_id}"
    try:
        res = requests.get(url, timeout=3).json()
        return res.get("details", "No detailed description available.")
    except:
        return "Could not retrieve malware details."

def osv_check(ecosystem, pkg_name, version):
    if interrupted: return {"flagged": False}
    url = "https://api.osv.dev/v1/query"
    eco_name = "PyPI" if ecosystem.lower() == "pypi" else "npm"
    try:
        # Check Name
        payload_name = {"package": {"name": pkg_name, "ecosystem": eco_name}}
        res_name = requests.post(url, json=payload_name, timeout=3).json()
        if "vulns" in res_name:
            for v in res_name["vulns"]:
                if v.get("id", "").startswith("MAL-"):
                    return {"flagged": True, "id": v["id"], "is_malware": True, "summary": v.get("summary"), "details": get_osv_details(v["id"])}
        
        # Check Version
        payload_ver = {"version": version, "package": {"name": pkg_name, "ecosystem": eco_name}}
        res_ver = requests.post(url, json=payload_ver, timeout=3).json()
        if "vulns" in res_ver:
            latest = res_ver["vulns"][0]
            return {"flagged": True, "id": latest.get("id"), "is_malware": False, "summary": latest.get("summary")}
    except: pass
    return {"flagged": False}

def osm_check(pkg_name, version, ecosystem, osm_key):
    if interrupted: return {"flagged": False}
    base_url = "https://api.opensourcemalware.com/functions/v1/check-malicious"
    headers = {"Authorization": f"Bearer {osm_key}"}
    params = {"report_type": "package", "resource_identifier": pkg_name, "version": version, "ecosystem": ecosystem.lower()}
    try:
        resp = requests.get(base_url, params=params, headers=headers, timeout=3)
        data = resp.json()
        if data.get("malicious"):
            return {"flagged": True, "reason": data.get("reason", "Malicious intent detected")}
    except: pass
    return {"flagged": False}

# --- Dependency Extraction ---

def get_dependencies(api_instance, pod_name, namespace, container_name):
    if interrupted: return []
    results = []
    # Try Python
    try:
        res_pip = stream(api_instance.connect_get_namespaced_pod_exec, pod_name, namespace, container=container_name, command=['pip', 'freeze'], stderr=False, stdin=False, stdout=True, tty=False)
        for line in res_pip.split('\n'):
            if '==' in line:
                name, ver = line.split('==')
                results.append(('PyPI', name.strip(), ver.strip()))
    except: pass
    
    if interrupted: return results

    # Try Node.js
    try:
        res_npm = stream(api_instance.connect_get_namespaced_pod_exec, pod_name, namespace, container=container_name, command=['npm', 'list', '--json', '--depth=0'], stderr=False, stdin=False, stdout=True, tty=False)
        npm_data = json.loads(res_npm)
        for name, info in npm_data.get('dependencies', {}).items():
            if info.get('version'):
                results.append(('npm', name, info['version']))
    except: pass
    return results

# --- Main Logic ---

def main():
    global interrupted
    parser = argparse.ArgumentParser(description="K8s Malware Scanner")
    parser.add_argument("-n", "--namespace", type=str, default="default", help="Namespace to scan")
    args = parser.parse_args()
    
    namespace = args.namespace
    osm_key = os.getenv("OSM_KEY")
    if not osm_key:
        print(f"{RED}Error: OSM_KEY not set.{RESET}"); sys.exit(1)

    try:
        config.load_kube_config()
    except:
        config.load_incluster_config()

    v1 = client.CoreV1Api()
    stats = Counter(total_scanned=0, clean=0, vulnerable=0, malicious=0)
    malicious_list = []

    print(f"\n{'='*75}\nKUBERNETES DEPENDENCY AUDIT\nTarget Namespace: {BOLD}{namespace}{RESET}\n{'='*75}")
    
    try:
        pods = v1.list_namespaced_pod(namespace)
        
        for pod in pods.items:
            if interrupted: break # Exit pod loop
            pod_name = pod.metadata.name
            
            for container in pod.spec.containers:
                if interrupted: break # Exit container loop
                print(f"\n[*] Pod: {pod_name} | Container: {container.name}")
                deps = get_dependencies(v1, pod_name, namespace, container.name)
                
                if not deps:
                    print("    [?] No dependencies detected.")
                    continue

                for eco, pkg, ver in deps:
                    if interrupted: break # Exit dependency loop
                    
                    stats['total_scanned'] += 1
                    osv_res = osv_check(eco, pkg, ver)
                    osm_res = osm_check(pkg, ver, eco, osm_key)

                    is_malware = osm_res["flagged"] or (osv_res.get("flagged") and osv_res.get("is_malware"))
                    
                    if is_malware:
                        reason = osm_res.get("reason") or osv_res.get("summary")
                        status = f"{RED}[MALWARE] {reason}{RESET}"
                        stats['malicious'] += 1
                        malicious_list.append({
                            "id": osv_res.get("id", "OSM-DETECT"),
                            "pkg": f"{pkg}@{ver}",
                            "eco": eco,
                            "pod": pod_name,
                            "details": osv_res.get("details", "OSM: " + osm_res.get("reason", "N/A"))
                        })
                    elif osv_res.get("flagged"):
                        status = f"{YELLOW}[VULN] {osv_res['id']}: {osv_res['summary']}{RESET}"
                        stats['vulnerable'] += 1
                    else:
                        status = f"{GREEN}[CLEAN]{RESET}"
                        stats['clean'] += 1

                    print(f"    - [{eco}] {pkg}@{ver}".ljust(60) + status)
    
    except Exception as e:
        if not interrupted:
            print(f"{RED}Error: {e}{RESET}")
            sys.exit(1)

    # --- Summary Report ---
    print(f"\n\n{'='*75}\nSCAN SUMMARY REPORT {'(PARTIAL)' if interrupted else '(COMPLETE)'}\n{'='*75}")
    print(f"Total Dependencies Scanned: {stats['total_scanned']}")
    print(f"Clean Packages:             {GREEN}{stats['clean']}{RESET}")
    print(f"Vulnerable Packages (CVE):  {YELLOW}{stats['vulnerable']}{RESET}")
    print(f"Malicious Packages:         {RED}{stats['malicious']}{RESET}")
    
    if malicious_list:
        print(f"\n\033[41m\033[97m !!! ACTION REQUIRED: MALICIOUS ENTRIES DETECTED !!! {RESET}")
        for entry in malicious_list:
            m_id = entry['id']
            display_id = f"{RED}{BOLD}{m_id}{RESET}" if m_id.startswith("MAL-") else m_id
            print(f"\n--- Threat Report: {entry['pkg']} ---")
            print(f"Location:  {entry['pod']} ({entry['eco']})")
            print(f"OSV ID:    {display_id}")
            if m_id.startswith("MAL-") or m_id.startswith("GHSA-"):
                print(f"OSV Link:  https://osv.dev/vulnerability/{m_id}")
            print(f"Analysis:  {entry['details']}")
            print("-" * 40)
    
    print(f"{'='*75}\nScan Finished.")

if __name__ == "__main__":
    main()