#!/usr/bin/env python3
"""CVE Scanner for WordPress plugins using Wordfence vulnerability database."""
import argparse
import csv
import json
import os
import sys


def load_json_local(path: str) -> dict:
    """Load JSON data from a local file.
    
    Args:
        path: Filesystem path to the JSON file.
        
    Returns:
        Parsed JSON data as a dictionary or list.
    """
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def load_json_from_url(url: str) -> dict:
    """Load JSON data from a URL.
    
    Args:
        url: URL to fetch JSON from.
        
    Returns:
        Parsed JSON data as a dictionary or list.
    """
    import requests
    r = requests.get(url, timeout=30)
    r.raise_for_status()
    return r.json()


def parse_args() -> argparse.Namespace:
    """Parse command-line arguments.
    
    Returns:
        Parsed arguments namespace.
    """
    parser = argparse.ArgumentParser(
        description="Scan WordPress plugins for CVE vulnerabilities using Wordfence data."
    )
    parser.add_argument(
        "--input",
        "-i",
        help="Path to local vulnerabilities JSON file (default: wordfence_vulnerabilities.json)",
        default="wordfence_vulnerabilities.json"
    )
    parser.add_argument(
        "--url",
        "-u",
        help="URL to download vulnerabilities JSON if local file not found",
        default=""
    )
    parser.add_argument(
        "--plugins",
        "-p",
        nargs="+",
        help="List of plugin slugs to check (space-separated)",
        default=[]
    )
    parser.add_argument(
        "--plugins-file",
        "-f",
        help="Path to file containing plugin slugs (one per line)",
        default=None
    )
    parser.add_argument(
        "--out-csv",
        "-c",
        help="Output CSV file path (if not specified, CSV output is skipped)",
        default=None
    )
    parser.add_argument(
        "--out-json",
        "-j",
        help="Output JSON file path (if not specified, prints to stdout)",
        default=None
    )
    return parser.parse_args()


def load_plugin_list(plugins: list[str], plugins_file: str | None) -> set[str]:
    """Load and combine plugin slugs from CLI args and optional file.
    
    Args:
        plugins: List of plugin slugs from command line.
        plugins_file: Optional path to file with plugin slugs (one per line).
        
    Returns:
        Set of plugin slugs to check.
    """
    plugin_set = set(plugins)
    
    if plugins_file and os.path.exists(plugins_file):
        with open(plugins_file, "r", encoding="utf-8") as f:
            for line in f:
                slug = line.strip()
                if slug and not slug.startswith("#"):
                    plugin_set.add(slug)
    
    return plugin_set


def main():
    """Main entry point for the CVE scanner."""
    args = parse_args()
    
    # Load plugin list
    bad_plugins = load_plugin_list(args.plugins, args.plugins_file)
    
    if not bad_plugins:
        print("Warning: No plugins specified to check. Use --plugins or --plugins-file.")
        print("Proceeding to load vulnerability data anyway...")
    
    # Load vulnerability data
    data = None
    # Prioritize URL if specified
    if args.url:
        print(f"Downloading from URL: {args.url}")
        data = load_json_from_url(args.url)
    elif os.path.exists(args.input):
        print(f"Loading local file: {args.input}")
        data = load_json_local(args.input)
    else:
        raise SystemExit(f"Error: Local file '{args.input}' not found and --url not provided.")

    # Extract vulnerability entries
    if isinstance(data, dict) and "vulnerabilities" in data:
        entries = data["vulnerabilities"]
    elif isinstance(data, list):
        entries = data
    else:
        entries = []
        if isinstance(data, dict):
            for v in data.values():
                if isinstance(v, list):
                    entries = v
                    break

    # Find matches
    matches = []
    for entry in entries:
        try:
            if entry.get("type") != "plugin":
                continue
            slug = entry.get("slug", "").strip()
            if not slug:
                continue
            if slug in bad_plugins:
                matches.append({
                    "slug": slug,
                    "title": entry.get("title"),
                    "cve": entry.get("cve"),
                    "severity": entry.get("severity"),
                    "cvss_score": entry.get("cvss_score"),
                    "cve_link": entry.get("cve_link"),
                    "raw": entry
                })
        except Exception:
            continue

    # Display results
    if not matches:
        print("No matches found for the specified plugins.")
    else:
        print(f"\nFound {len(matches)} vulnerability matches:")
        for m in matches:
            print(f"- {m['slug']} | CVE: {m.get('cve')} | sev: {m.get('severity')} | CVSS: {m.get('cvss_score')}")
            print(f"  link: {m.get('cve_link')}")

    # Write CSV output (only if specified)
    if args.out_csv:
        with open(args.out_csv, "w", newline='', encoding="utf-8") as csvf:
            writer = csv.DictWriter(csvf, fieldnames=["slug", "title", "cve", "severity", "cvss_score", "cve_link"])
            writer.writeheader()
            for m in matches:
                writer.writerow({
                    "slug": m["slug"],
                    "title": m.get("title"),
                    "cve": m.get("cve"),
                    "severity": m.get("severity"),
                    "cvss_score": m.get("cvss_score"),
                    "cve_link": m.get("cve_link")
                })
        print(f"\nSaved CSV results to {args.out_csv}")

    # Write JSON output (to file or stdout)
    if args.out_json:
        with open(args.out_json, "w", encoding="utf-8") as jf:
            json.dump(matches, jf, ensure_ascii=False, indent=2)
        print(f"Saved JSON results to {args.out_json}")
    else:
        # Print to stdout if no JSON file specified
        print("\nJSON output:")
        print(json.dumps(matches, ensure_ascii=False, indent=2))

if __name__ == "__main__":
    main()