from flask import Flask, request, jsonify
import json
import logging
import socket
from datetime import datetime, timezone

app = Flask(__name__)

# ======================================================
# Defense.com SYSLOG CEF2 Output Settings
# ======================================================
DEFENSE_SYSLOG_HOST = "x.x.x.x"   # <-- change
DEFENSE_SYSLOG_PORT = 5514                         # UDP 5514

# UDP socket
syslog_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)

# ======================================================
# Utility Functions (from your original code)
# ======================================================

def severity_to_cef(sev):
    """Map severity to CEF numeric."""
    if not sev:
        return 0
    sev = sev.upper()
    return {
        "LOW": 1,
        "MEDIUM": 3,
        "HIGH": 5,
        "CRITICAL": 8
    }.get(sev, 0)


def epoch_to_iso(epoch):
    """Convert epoch seconds to ISO."""
    try:
        return datetime.fromtimestamp(epoch, tz=timezone.utc).strftime(
            "%Y-%m-%dT%H:%M:%S.%fZ"
        )
    except Exception:
        return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ")


def extract_client_ip(value):
    if not value:
        return None
    if isinstance(value, list) and value:
        return value[0]
    if isinstance(value, str):
        return value
    return None


# ======================================================
# Build CEF2 + SYSLOG Message
# ======================================================

def build_cef_syslog(payload: dict) -> str:
    """
    Convert your incoming payload into a full SYSLOG + CEF2 line.
    """

    event_id = payload.get("id", "Unknown")
    severity = payload.get("severity", "Unknown")
    cef_sev = severity_to_cef(severity)
    detected_epoch = payload.get("detectedTime", 0)
    user = payload.get("userName", "Unknown")
    client_ip = extract_client_ip(payload.get("clientIPs"))
    protocol = payload.get("protocol", "Unknown")
    files = payload.get("files", [])
    num_files = payload.get("numFiles", len(files))
    state = payload.get("state", "Unknown")

    # CEF base header
    cef_header = (
        f"CEF:2|Superna|Data Security Edition|1.0|{event_id}|"
        f"{severity} severity event|{cef_sev}|"
    )

    # CEF extension fields
    cef_fields = [
        f"src={client_ip}" if client_ip else "",
        f"suser={user}" if user else "",
        f"protocol={protocol}" if protocol else "",
        f"cs1Label=State cs1={state}",
        f"cs2Label=NumFiles cs2={num_files}"
    ]

    # Add file paths as fname fields
    for f in files:
        cef_fields.append(f"fname={f}")

    cef_extension = " ".join([f for f in cef_fields if f])

    # RFC3164 syslog header
    now = datetime.now().strftime("%b %d %H:%M:%S")
    hostname = "superna"

    syslog_line = f"<134>{now} {hostname} superna-dse: {cef_header}{cef_extension}"

    return syslog_line


def send_cef_over_syslog(cef_message: str):
    """
    Send the CEF line to Defense.com collector via UDP.
    """
    # Print to console for debugging
    print("\n===== CEF MESSAGE SENT =====\n" + cef_message + "\n=============================\n")

    # Log as well
    logging.info(f"Sent CEF event to Defense.com collector: {cef_message}")

    # Send over UDP 5514
    syslog_socket.sendto(
        cef_message.encode('utf-8'),
        (DEFENSE_SYSLOG_HOST, DEFENSE_SYSLOG_PORT)
    )


# ======================================================
# Webhook Endpoint
# ======================================================

@app.route('/webhook', methods=['POST'])
def webhook():
    logging.debug("Incoming request payload: %s", request.data)

    payload = request.get_json(silent=True)
    if not payload:
        logging.error("Missing or invalid JSON payload")
        return jsonify({"error": "Missing or invalid JSON payload"}), 400

    # Build CEF syslog message
    cef_msg = build_cef_syslog(payload)

    # Send over UDP to Defense.com
    send_cef_over_syslog(cef_msg)

    return jsonify({"message": "Defense.com CEFv2 syslog event sent"}), 200


# ======================================================
# Main Entrypoint
# ======================================================

if __name__ == '__main__':
    logging.info("Starting Defense.com CEFv2 Syslog Webhook Listener on port 5000")
    app.run(host='0.0.0.0', port=5000, debug=True)