from flask import Flask, request
import logging
import json
from datetime import datetime, timezone
import requests
import sys

# === Flask App Initialization ===
app = Flask(__name__)

# === Cribl Configuration ===
# Use your Cribl Stream Cloud HTTP Bulk API input endpoint
CRIBL_URL = "https://default.main.goofy-shamir-iqst5ne.cribl.cloud:10080/cribl/_bulk"  # Use this for ECS format JSON payload
#CRIBL_URL = "https://default.main.goofy-shamir-iqst5ne.cribl.cloud:8088/services/collector/event"  # Use this for Splunk HEC format
CRIBL_TOKEN = "325a5246-3078-3057-396b-4b496d795738" # for  ecs


# === Format Toggle ===
# Set to True when sending to HEC input (port 8088)
# Set to False when sending to ECS /bulk input (port 10080)
hec_format = False

# === HTTP Headers ===
CRIBL_HEADERS = {
    "Content-Type": "application/json",
    "Authorization": f"Bearer {CRIBL_TOKEN}"
}

# === Logging Configuration ===
log_formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")

file_handler = logging.FileHandler("cribl.log")
file_handler.setFormatter(log_formatter)
file_handler.setLevel(logging.INFO)

console_handler = logging.StreamHandler(sys.stdout)
console_handler.setFormatter(log_formatter)
console_handler.setLevel(logging.INFO)

logger = logging.getLogger()
logger.setLevel(logging.INFO)
logger.handlers.clear()
logger.addHandler(file_handler)
logger.addHandler(console_handler)

# === ECS-Compatible Formatter ===
def format_for_cribl(payload):
    logger.info("🔧 Formatting payload (HEC=%s)", hec_format)

    device_product = "Data Security Edition"
    device_vendor = "Superna"
    event_type = "threat_detection"

    event_id = payload.get("id", "Unknown")
    severity = payload.get("severity", "Unknown")
    state = payload.get("state", "Unknown")
    detected = payload.get("detected", "Unknown")
    detected_time = payload.get("detectedTime", 0)

    try:
        epoch_time = detected_time / 1000
        timestamp = datetime.fromtimestamp(epoch_time, tz=timezone.utc).isoformat()
    except Exception as e:
        epoch_time = datetime.now(timezone.utc).timestamp()
        timestamp = datetime.fromtimestamp(epoch_time, tz=timezone.utc).isoformat()
        logger.error("❌ Error converting detectedTime: %s", str(e))

    client_ip = payload.get("clientIPs", ["Unknown"])[0]
    user = payload.get("userName", "Unknown")
    protocol = payload.get("protocol", "Unknown")
    url = payload.get("url", "Unknown")

    # Flattened ECS-style payload
    ecs_doc = {
        "@timestamp": timestamp,
        "event.id": event_id,
        "event.kind": "alert",
        "event.category": "threat",
        "event.type": event_type,
        "event.severity": severity,
        "event.state": state,
        "event.detected": detected,
        "event.user": user,
        "event.protocol": protocol,
        "event.client_ip": client_ip,
        "event.url": url,
        "observer.vendor": device_vendor,
        "observer.product": device_product
    }

    # === Convert to Splunk HEC format ===
    if hec_format:
        hec_doc = {
            "time": epoch_time,
            "host": "superna",
            "source": "superna:webhook",
            "sourcetype": "superna:alert",
            "event": ecs_doc
        }
        logger.info("✅ Formatted as Splunk HEC event:\n%s", json.dumps(hec_doc, indent=2))
        return hec_doc
    else:
        logger.info("✅ Formatted as ECS JSON:\n%s", json.dumps(ecs_doc, indent=2))
        return ecs_doc

# === Send to Cribl ===
def send_to_cribl(document):
    logger.info("📤 Sending document to Cribl")
    try:
        payload_json = json.dumps(document, indent=2)
        print("📄 Payload Sent:\n" + payload_json, flush=True)

        response = requests.post(
            CRIBL_URL,
            headers=CRIBL_HEADERS,
            data=payload_json
        )

        logger.info(f"📥 Cribl HTTP status: {response.status_code}")
        logger.info(f"📥 Cribl response: {response.text.strip()}")

        if response.status_code in [200, 201, 202]:
            try:
                cribl_response = response.json()
                cribl_count = cribl_response.get("count", None)
                logger.info(f"✅ Document successfully sent to Cribl (count={cribl_count})")

                # include record count in return message
                return {
                    "message": "Data sent to Cribl",
                    "cribl_status": response.status_code,
                    "record_count": cribl_count
                }
            except Exception:
                logger.warning("⚠️ Cribl returned non-JSON response")
                return {
                    "message": "Data sent to Cribl",
                    "cribl_status": response.status_code
                }
        else:
            logger.error("❌ Failed to send data to Cribl: %s", response.text)
            return {
                "error": f"Failed to send data: {response.text}",
                "cribl_status": response.status_code
            }
    except Exception as e:
        logger.exception("❌ Exception sending to Cribl")
        return {"error": str(e)}

# === Webhook Endpoint ===
@app.route('/webhook', methods=['POST'])
def webhook():
    logger.info("📩 Webhook endpoint triggered")
    try:
        payload = request.get_json(force=True)
        logger.info("📦 Incoming Payload:\n%s", json.dumps(payload, indent=2))
        print("📦 Raw Payload:\n" + json.dumps(payload, indent=2), flush=True)

        formatted_doc = format_for_cribl(payload)
        result = send_to_cribl(formatted_doc)

        logger.info("✅ Webhook processing completed")
        return json.dumps(result, indent=2), 200 if "message" in result else 500
    except Exception as e:
        logger.exception("❌ Error in webhook handler")
        return json.dumps({"error": str(e)}), 500

# === Run Server ===
if __name__ == '__main__':
    logger.info("🚀 Starting Cribl webhook server at 0.0.0.0:5000")
    app.run(host='0.0.0.0', port=5000, debug=True)