from flask import Flask, request
import logging
import json
from datetime import datetime, timezone
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
import requests
import gzip
from io import BytesIO

# Flask app initialization
app = Flask(__name__)

# HTTP Collector Endpoint and API Token
exabeam_URL = "https://api2.us-east.exabeam.cloud/cloud-collectors/v1/logs/json"
exabeam_api_token = "eyJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiJ9.eyJzdWJzY3JpcHRpb25Db2RlIjoiOWwzMzJ6IiwiYXVkIjoibmdjYy13ZWJob29rLWNhbGxiYWNrLTB3a3RvaGN1a3Nic2EuYXBpZ2F0ZXdheS5leGEtY2xvdWQtcHJvZC5jbG91ZC5nb29nIiwic3ViIjoiZXIyNmh0a3ZxbXExa2JpdDY2YzU1anY5dDFqdWd0QGV4YS1jbG91ZC1wcm9kLmlhbS5nc2VydmljZWFjY291bnQuY29tIiwiaW5zdGFuY2VJZCI6IjlkYmFlOTU4LTQzMTgtNDdiNC05ZWY0LTdmMmUzNDA3YzBiNCIsImxvZ0Zvcm1hdCI6Impzb24iLCJjb25uZWN0b3JJZCI6ImdlbmVyaWMtZXZlbnRzIiwiaXNzIjoiZXIyNmh0a3ZxbXExa2JpdDY2YzU1anY5dDFqdWd0QGV4YS1jbG91ZC1wcm9kLmlhbS5nc2VydmljZWFjY291bnQuY29tIiwiZW1haWwiOiJlcjI2aHRrdnFtcTFrYml0NjZjNTVqdjl0MWp1Z3RAZXhhLWNsb3VkLXByb2QuaWFtLmdzZXJ2aWNlYWNjb3VudC5jb20ifQ.kx3cToh_A-0TWvU0tTo0gEJp3I6dOqdYlANAwgUqxZ2rRWpvGgEkvu98DDcqe7874_zflZjCaJY0KH3Zsp4B6XHZ87rnMxJ4z5Q3zy5phz9yYfkRKd7lLh1Iwm_w-MSYxuPOLFGxo6FEBK2sJBsifl5spiOpYIkMyE7eGZ-nbhlMEF9xg-dMCopg5wGvvDUhLR5oGOLyIyAm6XHVBeX8qEAeCIj0PnQs4StN0Q01CdNruFagzh7TABCUyB6qzRIBtM4Bhpy75NOetMwlbySa8qkzo_5AvqXAdNNcU4drLg8hE6L8xsqTaN0sDcj39s-G3oVybXkSXWJ6wsbPx33tiQ"

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

# UEF Formatting Function
def format_uef(payload):
    # Static UEF fields
    device_product = "Eyeglass Zero Trust"
    device_vendor = "Superna"
    device_version = "V1"
    device_event_class_id = "security"
    version = "1.0"
    event_type = "threat_detection"

    # Extract dynamic fields from the payload
    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", []))

    # Convert detectedTime to ISO 8601 UTC format for the timestamp field
    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)

    # Format UEF message
    uef_message = {
        "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,
        "protocol": protocol,
        "client_ip": client_ip,
        "user": user,
        "action": action
    }

    return uef_message

# Utilities
def get_requests_session():
    session = requests.Session()
    retries = Retry(total=3, backoff_factor=1, status_forcelist=[500, 502, 503, 504])
    session.mount("https://", HTTPAdapter(max_retries=retries))
    return session

def send_to_siem(payload):
    headers = {
        "Content-Type": "application/json",
        "Content-Encoding": "gzip",
        "Authorization": f"Bearer {exabeam_api_token}"  # Add Bearer token
    }
    session = get_requests_session()

    try:
        # Console log for raw JSON payload before compression
        print("Raw JSON payload before compression:")
        print(json.dumps(payload, indent=4))

        # Compress the payload using gzip
        buffer = BytesIO()
        with gzip.GzipFile(fileobj=buffer, mode="w") as gz:
            gz.write(json.dumps(payload).encode("utf-8"))
        gzipped_payload = buffer.getvalue()

        # Console logs for payload compression and endpoint
        print("Payload successfully compressed.")
        print(f"Sending payload to URL: {exabeam_URL}")

        # Logging the payload for debugging
        logging.info("Payload Sent: %s", json.dumps(payload))

        # Send gzipped payload
        response = session.post(exabeam_URL, headers=headers, data=gzipped_payload)
        if response.status_code in [200, 204]:
            print(f"Success: Status Code {response.status_code}")
            logging.info("Success: Status Code %s", response.status_code)
        else:
            print(f"Error: Status Code {response.status_code}, Response: {response.text}")
            logging.error("Failed: Status Code %s, Response: %s", response.status_code, response.text)
    except Exception as e:
        print(f"Exception: {str(e)}")
        logging.error("Exception: %s", str(e))

# Flask webhook route
@app.route('/webhook', methods=['POST'])
def webhook():
    try:
        payload = request.json
        # Format the UEF message
        uef_message = format_uef(payload)
        # Send UEF message to SIEM
        send_to_siem(uef_message)
        return "Success", 200
    except Exception as e:
        logging.error("Error handling webhook: %s", str(e))
        return str(e), 500

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