app.py

import re
import os
from flask import Flask, render_template, request, jsonify, Response
from scanner import scan_target
from cve_lookup import lookup_cves_for_services
from database import init_db, save_scan, get_all_scans, get_cves_by_scan, get_severity_summary, get_latest_scan_id, get_all_cves
from reporter import generate_csv, generate_pdf

app = Flask(__name__)

def is_valid_target(target: str) -> bool:
    """Validate if target is a valid IPv4 or domain name (including localhost)."""
    target = target.strip()
    if target.lower() == "localhost":
        return True
    ipv4_pattern = r"^(?:(?:25[0-5]|2[0-4][0-9]|[0-1]?[0-9]{1,2})\.){3}(?:25[0-5]|2[0-4][0-9]|[0-1]?[0-9]{1,2})$"
    hostname_pattern = r"^(?:[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?\.)+[a-zA-Z]{2,6}$"
    return bool(re.match(ipv4_pattern, target) or re.match(hostname_pattern, target))

# Initialize database on startup
init_db()


@app.route("/")
def index():
    """Main dashboard — shows all past scans."""
    scans = get_all_scans()
    latest_id = get_latest_scan_id()

    severity_summary = {"CRITICAL": 0, "HIGH": 0, "MEDIUM": 0, "LOW": 0, "N/A": 0}
    latest_cves = []

    if latest_id:
        severity_summary = get_severity_summary(latest_id)
        latest_cves = get_cves_by_scan(latest_id)

    nvd_api_key = os.getenv("NVD_API_KEY", "")

    return render_template(
        "index.html",
        scans=scans,
        latest_cves=latest_cves,
        severity_summary=severity_summary,
        latest_id=latest_id,
        nvd_api_key=nvd_api_key
    )


@app.route("/scan", methods=["POST"])
def run_scan():
    """
    Accepts JSON: { "target": "127.0.0.1" }
    Runs Nmap scan → CVE lookup → saves to DB.
    Returns scan_id and summary.
    """
    data = request.get_json()

    if not data or "target" not in data:
        return jsonify({"error": "No target provided"}), 400

    target = data["target"].strip()

    if not target:
        return jsonify({"error": "Target cannot be empty"}), 400

    if not is_valid_target(target):
        return jsonify({"error": "Invalid target format. Must be a valid IP address or domain name."}), 400

    print(f"\n[*] Scan requested for: {target}")

    # Step 1: Nmap scan
    scan_data = scan_target(target)

    if not scan_data or not scan_data.get("services"):
        return jsonify({
            "error": f"No open services found on {target}. Host may be down or unreachable."
        }), 404

    # Step 2: CVE lookup for each service
    cve_results = lookup_cves_for_services(scan_data["services"])

    # Step 3: Save to database
    scan_id = save_scan(scan_data, cve_results)

    return jsonify({
        "scan_id": scan_id,
        "target": target,
        "total_services": len(scan_data["services"]),
        "total_cves": len(cve_results),
        "message": f"Scan complete. Found {len(cve_results)} CVE(s) across {len(scan_data['services'])} service(s)."
    })


@app.route("/results/<int:scan_id>")
def get_results(scan_id):
    """
    Returns CVE results and severity summary for a specific scan as JSON.
    Used by the dashboard to load results dynamically.
    """
    cves = get_cves_by_scan(scan_id)
    summary = get_severity_summary(scan_id)
    scans = get_all_scans()

    # Find the scan record
    scan = next((s for s in scans if s["id"] == scan_id), None)

    if not scan:
        return jsonify({"error": "Scan not found"}), 404

    return jsonify({
        "scan": scan,
        "cves": cves,
        "severity_summary": summary
    })


@app.route("/scans")
def list_scans():
    """Returns all past scans as JSON."""
    scans = get_all_scans()
    return jsonify(scans)


@app.route("/export/csv/<int:scan_id>")
def export_csv(scan_id):
    """Generate and download a CSV report for a given scan."""
    scans = get_all_scans()
    scan = next((s for s in scans if s["id"] == scan_id), None)
    if not scan:
        return jsonify({"error": "Scan not found"}), 404

    cves = get_cves_by_scan(scan_id)
    csv_content = generate_csv(scan, cves)

    filename = f"cve_report_{scan['target']}_{scan_id}.csv"
    return Response(
        csv_content,
        mimetype="text/csv",
        headers={"Content-disposition": f"attachment; filename={filename}"}
    )


@app.route("/export/pdf/<int:scan_id>")
def export_pdf(scan_id):
    """Generate and download a PDF report for a given scan."""
    scans = get_all_scans()
    scan = next((s for s in scans if s["id"] == scan_id), None)
    if not scan:
        return jsonify({"error": "Scan not found"}), 404

    cves = get_cves_by_scan(scan_id)
    pdf_bytes = generate_pdf(scan, cves)

    filename = f"cve_report_{scan['target']}_{scan_id}.pdf"
    return Response(
        pdf_bytes,
        mimetype="application/pdf",
        headers={"Content-disposition": f"attachment; filename={filename}"}
    )


@app.route("/api/cves")
def list_all_cves():
    """Returns all CVE results across all scans as JSON."""
    cves = get_all_cves()
    return jsonify(cves)


@app.route("/api/settings", methods=["POST"])
def update_settings():
    """Accepts JSON: { "api_key": "..." } and updates NVD_API_KEY in .env."""
    data = request.get_json()
    if not data or "api_key" not in data:
        return jsonify({"error": "No api_key provided"}), 400

    api_key = data["api_key"].strip()

    # Read existing env file if present
    env_lines = []
    if os.path.exists(".env"):
        try:
            with open(".env", "r") as f:
                env_lines = f.readlines()
        except Exception:
            pass

    # Update or add NVD_API_KEY
    key_found = False
    new_lines = []
    for line in env_lines:
        if line.strip().startswith("NVD_API_KEY="):
            new_lines.append(f"NVD_API_KEY={api_key}\n")
            key_found = True
        else:
            new_lines.append(line)

    if not key_found:
        new_lines.append(f"NVD_API_KEY={api_key}\n")

    try:
        with open(".env", "w") as f:
            f.writelines(new_lines)
    except Exception as e:
        return jsonify({"error": f"Failed to write to .env: {e}"}), 500

    # Reload environment variable locally so the running server uses it immediately
    os.environ["NVD_API_KEY"] = api_key

    return jsonify({"success": True, "message": "API key saved successfully."})


if __name__ == "__main__":
    debug_mode = os.getenv("FLASK_DEBUG", "false").lower() in ("true", "1")
    app.run(debug=debug_mode, port=5000)