app.py

"""
app.py
------
Flask entry point for the Vulnerability Management dashboard.
Defines all REST API routes and serves the single-page frontend.
"""

from flask import Flask, request, jsonify, render_template
from database import init_db, get_db

app = Flask(__name__)


# ---------------------------------------------------------------------------
# Serve the dashboard SPA
# ---------------------------------------------------------------------------

@app.route("/")
def index():
    """Render the main dashboard HTML page."""
    return render_template("index.html")


# ---------------------------------------------------------------------------
# Stats endpoint – used to populate KPI cards
# ---------------------------------------------------------------------------

@app.route("/api/stats")
def get_stats():
    """Return vulnerability counts grouped by severity and overall totals."""
    conn = get_db()
    cursor = conn.cursor()

    # Total count
    total = cursor.execute("SELECT COUNT(*) FROM vulnerabilities").fetchone()[0]

    # Count per severity
    severity_rows = cursor.execute(
        "SELECT severity, COUNT(*) as cnt FROM vulnerabilities GROUP BY severity"
    ).fetchall()

    # Count per status
    status_rows = cursor.execute(
        "SELECT status, COUNT(*) as cnt FROM vulnerabilities GROUP BY status"
    ).fetchall()

    conn.close()

    severity_counts = {row["severity"]: row["cnt"] for row in severity_rows}
    status_counts   = {row["status"]:   row["cnt"] for row in status_rows}

    return jsonify({
        "total":    total,
        "critical": severity_counts.get("Critical", 0),
        "high":     severity_counts.get("High", 0),
        "medium":   severity_counts.get("Medium", 0),
        "low":      severity_counts.get("Low", 0),
        "open":       status_counts.get("Open", 0),
        "in_progress": status_counts.get("In Progress", 0),
        "resolved":    status_counts.get("Resolved", 0),
    })


# ---------------------------------------------------------------------------
# Vulnerabilities – CRUD routes
# ---------------------------------------------------------------------------

@app.route("/api/vulnerabilities", methods=["GET"])
def list_vulnerabilities():
    """
    Return all vulnerabilities, optionally filtered by severity or status.
    Query params: ?severity=Critical&status=Open
    """
    conn = get_db()
    cursor = conn.cursor()

    query  = "SELECT * FROM vulnerabilities WHERE 1=1"
    params = []

    severity = request.args.get("severity")
    if severity:
        query += " AND severity = ?"
        params.append(severity)

    status = request.args.get("status")
    if status:
        query += " AND status = ?"
        params.append(status)

    query += " ORDER BY CASE severity WHEN 'Critical' THEN 1 WHEN 'High' THEN 2 WHEN 'Medium' THEN 3 WHEN 'Low' THEN 4 END, id DESC"

    rows = cursor.execute(query, params).fetchall()
    conn.close()

    return jsonify([dict(row) for row in rows])


@app.route("/api/vulnerabilities", methods=["POST"])
def create_vulnerability():
    """Create a new vulnerability record from JSON body."""
    data = request.get_json()
    if not data:
        return jsonify({"error": "No data provided"}), 400

    required = ["title", "severity"]
    missing = [f for f in required if not data.get(f)]
    if missing:
        return jsonify({"error": f"Missing required fields: {', '.join(missing)}"}), 400

    conn = get_db()
    cursor = conn.cursor()
    cursor.execute("""
        INSERT INTO vulnerabilities (title, description, severity, status, affected_asset, discovered_date)
        VALUES (?, ?, ?, ?, ?, ?)
    """, (
        data.get("title"),
        data.get("description", ""),
        data.get("severity"),
        data.get("status", "Open"),
        data.get("affected_asset", ""),
        data.get("discovered_date", ""),
    ))
    conn.commit()
    new_id = cursor.lastrowid
    new_row = cursor.execute("SELECT * FROM vulnerabilities WHERE id = ?", (new_id,)).fetchone()
    conn.close()

    return jsonify(dict(new_row)), 201


@app.route("/api/vulnerabilities/<int:vuln_id>", methods=["PUT"])
def update_vulnerability(vuln_id):
    """Update an existing vulnerability by ID."""
    data = request.get_json()
    if not data:
        return jsonify({"error": "No data provided"}), 400

    conn = get_db()
    cursor = conn.cursor()

    # Verify it exists
    existing = cursor.execute("SELECT * FROM vulnerabilities WHERE id = ?", (vuln_id,)).fetchone()
    if not existing:
        conn.close()
        return jsonify({"error": "Vulnerability not found"}), 404

    cursor.execute("""
        UPDATE vulnerabilities
        SET title = ?, description = ?, severity = ?, status = ?, affected_asset = ?, discovered_date = ?
        WHERE id = ?
    """, (
        data.get("title",          existing["title"]),
        data.get("description",    existing["description"]),
        data.get("severity",       existing["severity"]),
        data.get("status",         existing["status"]),
        data.get("affected_asset", existing["affected_asset"]),
        data.get("discovered_date",existing["discovered_date"]),
        vuln_id,
    ))
    conn.commit()
    updated = cursor.execute("SELECT * FROM vulnerabilities WHERE id = ?", (vuln_id,)).fetchone()
    conn.close()

    return jsonify(dict(updated))


@app.route("/api/vulnerabilities/<int:vuln_id>", methods=["DELETE"])
def delete_vulnerability(vuln_id):
    """Delete a vulnerability by ID."""
    conn = get_db()
    cursor = conn.cursor()

    existing = cursor.execute("SELECT id FROM vulnerabilities WHERE id = ?", (vuln_id,)).fetchone()
    if not existing:
        conn.close()
        return jsonify({"error": "Vulnerability not found"}), 404

    cursor.execute("DELETE FROM vulnerabilities WHERE id = ?", (vuln_id,))
    conn.commit()
    conn.close()

    return jsonify({"message": "Vulnerability deleted successfully"}), 200


# ---------------------------------------------------------------------------
# App entry point
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    # Initialize the database (creates tables and seeds data if needed)
    init_db()
    app.run(debug=True, port=5000)