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

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

# === TheHive Configuration ===
THEHIVE_URL = "http://x.x.x.x:9000/api"
THEHIVE_API_KEY = "yyyyyyyy"
THEHIVE_HEADERS = {
    "Authorization": f"Bearer {THEHIVE_API_KEY}",
    "Content-Type": "application/json"
}

# === Logging Configuration ===
logging.basicConfig(
    filename="app.log",
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
)

def map_severity(severity_str):
    severity_map = {
        "critical": 3,
        "major": 2,
        "warning": 1
    }
    return severity_map.get(str(severity_str).lower(), 2)

def format_uef(payload):
    device_product = "Data Security Edition"
    device_vendor = "Superna"
    device_version = "V1"
    device_event_class_id = "security"
    version = "1.0"
    event_type = "threat_detection"

    event = 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)
    nes = ", ".join(payload.get("nes", []))
    files = "; ".join(payload.get("files", []))
    shares = "; ".join(share.get("name", "Unknown") for share in payload.get("shares", []))
    alert_url = payload.get("url", "Unknown")

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

    protocol = payload.get("protocol", "Unknown")
    client_ip = payload.get("clientIPs", ["Unknown"])[0]
    user = payload.get("userName", "Unknown")
    actions = [action.get("action", "Unknown") for action in payload.get("actions", [])]
    action = " | ".join(actions)

    return {
        "timestamp": timestamp,
        "device_product": device_product,
        "device_vendor": device_vendor,
        "device_version": device_version,
        "device_event_class_id": device_event_class_id,
        "version": version,
        "event_type": event_type,
        "event": event,
        "severity": severity,
        "state": state,
        "detected": detected,
        "detected_time": detected_time,
        "nes": nes,
        "files": files,
        "shares": shares,
        "protocol": protocol,
        "client_ip": client_ip,
        "user": user,
        "action": action,
        "Alert_url": alert_url
    }

def search_case_by_event_id(event_id):
    search_url = f"{THEHIVE_URL}/case/_search"
    query = {
        "filter": {
            "_string": f'title:"Incident id {event_id}"'
        }
    }

    print(f"🔍 Searching for case with event ID: {event_id}")
    print(f"➡️ Search URL: {search_url}")
    print(f"➡️ Query body: {json.dumps(query, indent=2)}")

    try:
        response = requests.post(search_url, headers=THEHIVE_HEADERS, data=json.dumps(query))
        print(f"⬅️ Search Response: {response.status_code} {response.text}")
        if response.status_code == 200:
            cases = response.json()
            return cases[0] if cases else None
        else:
            logging.error("Case search failed: %s - %s", response.status_code, response.text)
    except Exception as e:
        logging.error("Exception during case search: %s", str(e))
    return None

def add_case_comment(case_id, message):
    url = f"{THEHIVE_URL}/v1/case/{case_id}/comment"
    payload = {
        "message": message,
        "_type": "case"
    }

    print(f"📘 Adding comment to case ID: {case_id}")
    print("➡️ Comment URL:", url)
    print("➡️ Payload:", json.dumps(payload, indent=2))

    try:
        response = requests.post(url, headers=THEHIVE_HEADERS, data=json.dumps(payload))
        print(f"⬅️ Comment Response: {response.status_code} {response.text}")
        if response.status_code == 201:
            print("✅ Comment added successfully.")
        else:
            print(f"❌ Failed to add comment: {response.status_code} - {response.text}")
            logging.error("Failed to add comment: %s - %s", response.status_code, response.text)
    except Exception as e:
        print(f"❌ Exception adding comment: {str(e)}")
        logging.error("Exception adding comment: %s", str(e))

def create_thehive_case_from_uef(uef):
    event_id = uef.get("event", "Unknown")
    title = f"Superna zero trust: Suspicious user behavior for user {uef.get('user', 'Unknown')} - Incident id {event_id}"
    severity_value = map_severity(uef.get("severity"))

    case_data = {
        "title": title,
        "description": (
            f"Detected: {uef.get('detected')}\n"
            f"Timestamp: {uef.get('timestamp')}\n"
            f"User: {uef.get('user')}\n"
            f"Client IP: {uef.get('client_ip')}\n"
            f"Protocol: {uef.get('protocol')}\n"
            f"Files: {uef.get('files')}\n"
            f"Shares: {uef.get('shares')}\n"
            f"Action: {uef.get('action')}\n"
            f"Alert URL: {uef.get('Alert_url')}\n"
        ),
        "severity": severity_value,
        "tlp": 2,
        "pap": 2,
        "tags": ["superna", "webhook", uef.get("event_type", "alert")]
    }

    print("📤 Creating new case in TheHive...")
    print("➡️ Payload:", json.dumps(case_data, indent=2))

    try:
        response = requests.post(f"{THEHIVE_URL}/case", headers=THEHIVE_HEADERS, data=json.dumps(case_data))
        print("⬅️ Response Code:", response.status_code)
        print("⬅️ Response Body:", response.text)
        if response.status_code == 201:
            case_id = response.json().get("id")
            print(f"✅ TheHive case created: {case_id}")
            return {"thehive_case_id": case_id}
        else:
            logging.error("Failed to create TheHive case: %s - %s", response.status_code, response.text)
            return {"error": f"Failed to create case: {response.text}"}
    except Exception as e:
        logging.error("Exception creating TheHive case: %s", str(e))
        return {"error": str(e)}

@app.route('/webhook', methods=['POST'])
def webhook():
    try:
        payload = request.json
        print("📩 Webhook received:")
        print(json.dumps(payload, indent=2))
        logging.info("Webhook received: %s", json.dumps(payload))

        event_id = payload.get("id", "Unknown")
        existing_case = search_case_by_event_id(event_id)

        if existing_case:
            print(f"✅ Found existing case: {existing_case.get('title')}")
            case_id = existing_case.get("id")
            uef_message = format_uef(payload)
            comment = (
                f"📩 Webhook update received\n"
                f"Detected: {uef_message.get('detected')}\n"
                f"Timestamp: {uef_message.get('timestamp')}\n"
                f"User: {uef_message.get('user')}\n"
                f"Client IP: {uef_message.get('client_ip')}\n"
                f"Protocol: {uef_message.get('protocol')}\n"
                f"Files: {uef_message.get('files')}\n"
                f"Shares: {uef_message.get('shares')}\n"
                f"Action: {uef_message.get('action')}\n"
                f"Alert URL: {uef_message.get('Alert_url')}\n"
            )
            add_case_comment(case_id, comment)
            return json.dumps({"message": "Case updated"}), 200
        else:
            uef_message = format_uef(payload)
            return json.dumps(create_thehive_case_from_uef(uef_message)), 200
    except Exception as e:
        print(f"❌ Error handling webhook: {str(e)}")
        logging.error("Error handling webhook: %s", str(e))
        return json.dumps({"error": str(e)}), 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=True)