from flask import Flask, request
import argparse
import os
import json
import time
import logging
from datetime import datetime
import requests
from cryptography.fernet import Fernet

# Configure logging
logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(levelname)s - %(message)s'
)

app = Flask(__name__)

# Define constants
TOKEN_URL = "https://accounts.zohocloud.ca/oauth/v2/token"
API_URL = "https://servicedeskplus.ca/api/v3/problems"
NOTES_API_URL = "https://servicedeskplus.ca/api/v3/problems/{problem_id}/notes"
TOKEN_FILE = "tokens.enc"
KEY_FILE = "secret.key"
SELF_CLIENT_FILE = "self_client.json"

# Global variables
access_token = None
refresh_token = None
token_expiry = 0

# Encryption setup
def load_or_generate_key():
    if not os.path.exists(KEY_FILE):
        key = Fernet.generate_key()
        with open(KEY_FILE, "wb") as f:
            f.write(key)
    else:
        with open(KEY_FILE, "rb") as f:
            key = f.read()
    return Fernet(key)

cipher = load_or_generate_key()

def save_tokens_to_file(data):
    encrypted_data = cipher.encrypt(json.dumps(data).encode())
    with open(TOKEN_FILE, "wb") as f:
        f.write(encrypted_data)

def load_tokens_from_file():
    if not os.path.exists(TOKEN_FILE):
        raise FileNotFoundError("Token file not found. Please generate tokens.")
    with open(TOKEN_FILE, "rb") as f:
        encrypted_data = f.read()
    return json.loads(cipher.decrypt(encrypted_data).decode())

def load_client_config():
    file_path = os.path.join(os.getcwd(), SELF_CLIENT_FILE)
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"Client config file '{SELF_CLIENT_FILE}' not found.")
    with open(file_path, "r") as f:
        return json.load(f)

def get_access_token(client_id, client_secret, authorization_code, redirect_uri):
    global access_token, refresh_token, token_expiry

    payload = {
        "grant_type": "authorization_code",
        "code": authorization_code,
        "client_id": client_id,
        "client_secret": client_secret,
        "redirect_uri": redirect_uri,
    }

    response = requests.post(TOKEN_URL, data=payload)
    if response.status_code == 200:
        token_data = response.json()
        access_token = token_data.get("access_token")
        refresh_token = token_data.get("refresh_token")
        token_expiry = time.time() + token_data.get("expires_in", 0)

        save_tokens_to_file({
            "refresh_token": refresh_token,
            "access_token": access_token,
            "token_expiry": token_expiry,
        })
        logging.info("Tokens saved successfully.")
    else:
        raise Exception(f"Authorization failed. Response: {response.text}")

def refresh_access_token(client_id, client_secret):
    global access_token, token_expiry
    tokens = load_tokens_from_file()
    refresh_token = tokens["refresh_token"]

    payload = {
        "grant_type": "refresh_token",
        "refresh_token": refresh_token,
        "client_id": client_id,
        "client_secret": client_secret,
    }

    response = requests.post(TOKEN_URL, data=payload)
    if response.status_code == 200:
        token_data = response.json()
        access_token = token_data.get("access_token")
        token_expiry = time.time() + token_data.get("expires_in", 0)

        save_tokens_to_file({
            "refresh_token": refresh_token,
            "access_token": access_token,
            "token_expiry": token_expiry,
        })

        expiry_time_human = datetime.utcfromtimestamp(token_expiry).strftime('%Y-%m-%d %H:%M:%S UTC')
        logging.info(f"Access token refreshed successfully. New expiry time: {expiry_time_human}")
    else:
        logging.error(f"Token refresh failed. Response: {response.text}")
        raise Exception("Token refresh failed.")

def ensure_access_token(client_id, client_secret):
    global access_token, token_expiry
    try:
        tokens = load_tokens_from_file()
        access_token = tokens["access_token"]
        token_expiry = tokens["token_expiry"]
        if time.time() >= token_expiry:
            logging.info("Access token expired. Refreshing token...")
            refresh_access_token(client_id, client_secret)
    except FileNotFoundError:
        raise Exception("Token file not found. Please generate tokens.")

def format_udm_to_html(payload, source_ip):
    """Format the UDM payload into HTML format."""
    share_names = [share["name"] for share in payload.get("shares", []) if "name" in share]
    udm_event = {
        "incidentmetadata": {
            "Event ID": payload.get('id'),
            "event_type": "Cyber Storage",
            "event_timestamp": datetime.utcnow().strftime('%Y-%m-%dT%H:%M:%SZ'),
            "product_log_id": "Superna Zero Trust Webhook",
            "vendor_name": "Superna",
            "product_name": "Eyeglass Data Security Edition",
            "product_version": "V2.5.9",
        },
        "principal": {
            "hostname": payload.get('clientIPs', ['Unknown'])[0],
            "ip": payload.get('clientIPs', ['Unknown'])[0],
            "user": {"user_id": payload.get('userName', 'Unknown')},
        },
        "target": {
            "hostname": payload.get('clientIPs', ['Unknown'])[0],
            "ip": payload.get('clientIPs', ['Unknown'])[0],
        },
        "security_result": {
            "about": "Superna Zero Trust Cyber Storage Threat Detection",
            "category": "SOFTWARE_MALICIOUS",
            "severity": payload.get('severity', 'Unknown'),
            "target_of_attack": payload.get('nes', []),
            "state": payload.get('state'),
            "share_names": share_names,
        },
        "file_list": payload.get('files', []),
    }
    return "<pre>" + json.dumps(udm_event, indent=4) + "</pre>"

def send_to_service_desk(title, description, priority_name, status_name, impact_details, root_cause):
    ensure_access_token(CLIENT_ID, CLIENT_SECRET)

    problem_data = {
        "problem": {
            "title": title,
            "description": description,
            "priority": {"name": priority_name},
            "status": {"name": status_name},
            "impact_details": {"impact_details_description": impact_details},
            "root_cause": {"root_cause_description": root_cause},
        }
    }

    headers = {
        "Authorization": f"Zoho-oauthtoken {access_token}",
        "Content-Type": "application/x-www-form-urlencoded",
        "Accept": "application/vnd.manageengine.sdp.v3+json",
    }

    payload = {"input_data": json.dumps(problem_data)}
    response = requests.post(API_URL, headers=headers, data=payload)

    if response.status_code == 201:
        response_json = response.json()
        problem_id = response_json["problem"]["id"]
        logging.info(f"Problem created successfully. Problem ID: {problem_id}")
        return problem_id
    else:
        logging.error(f"Failed to create problem. Response: {response.text}")
        return None

def update_problem_notes(problem_id, notes, notify_users=None, notify_xpaths=None):
    """
    Update problem notes using the ManageEngine API.

    :param problem_id: ID of the problem to update notes for.
    :param notes: The note content as a string.
    :param notify_users: List of users to notify (default: None).
    :param notify_xpaths: List of xpaths to notify (default: None).
    """
    ensure_access_token(CLIENT_ID, CLIENT_SECRET)

    url = NOTES_API_URL.format(problem_id=problem_id)
    headers = {
        "Authorization": f"Zoho-oauthtoken {access_token}",
        "Content-Type": "application/x-www-form-urlencoded",
        "Accept": "application/vnd.manageengine.sdp.v3+json",
    }

    # Construct the payload based on the API format
    note_payload = {
        "note": {
            "description": notes,
        }
    }

    payload = {"input_data": json.dumps(note_payload)}
    response = requests.post(url, headers=headers, data=payload)

    if response.status_code == 201:
        logging.info(f"Notes updated successfully for Problem ID: {problem_id}")
    else:
        logging.error(f"Failed to update notes for Problem ID: {problem_id}. Response: {response.text}")

def extract_actions(payload):
    """Extract and format actions from the payload."""
    actions = payload.get("actions", [])
    if not actions:
        return "No actions found in the payload."

    action_notes = []
    for action in actions:
        formatted_action = (
            f"Action: {action.get('action', 'Unknown')}\n"
            f"Date: {datetime.utcfromtimestamp(action.get('dateInLong', 0) / 1000).strftime('%Y-%m-%d %H:%M:%S')}\n"
            f"Result State: {action.get('resultState', 'Unknown')}\n"
            f"Comment: {action.get('comment', 'No comment provided')}\n"
        )
        action_notes.append(formatted_action)
    return "\n".join(action_notes)

def fetch_existing_problems():
    """
    Fetch all existing problems from ManageEngine ServiceDesk Plus API.

    :return: List of problems or an empty list on failure.
    """
    ensure_access_token(CLIENT_ID, CLIENT_SECRET)

    headers = {
        "Authorization": f"Zoho-oauthtoken {access_token}",
        "Accept": "application/vnd.manageengine.sdp.v3+json",
    }

    response = requests.get(API_URL, headers=headers)

    if response.status_code == 200:
        problems = response.json().get("problems", [])
        logging.info(f"Retrieved {len(problems)} problems from ServiceDesk.")
        return problems
    else:
        logging.error(f"Failed to fetch problems. Response: {response.text}")
        return []

def find_existing_problem(problems, event_id):
    """
    Check if a problem with the given event ID exists and is not closed.

    :param problems: List of existing problems.
    :param event_id: Event ID to search for in the title.
    :return: The problem ID if a matching problem is found; otherwise, None.
    """
    for problem in problems:
        title = problem.get("title", "")
        status = problem.get("status", {}).get("name", "").lower()
        if event_id in title and status != "closed":
            return problem.get("id")
    return None


@app.route('/webhook', methods=['POST'])
def webhook():
    try:
        payload = request.json
        source_ip = request.remote_addr
        description = format_udm_to_html(payload, source_ip)
        logging.info(f"Webhook received from {source_ip}.")

        event_id = payload.get('id', 'Unknown')
        title = f"Superna CyberStorage Zero Trust EventID {event_id} for user: {payload.get('userName', 'Unknown')}"

        # Fetch existing problems and check for duplicates
        existing_problems = fetch_existing_problems()
        problem_id = find_existing_problem(existing_problems, event_id)

        if problem_id:
            logging.info(f"Existing problem found with ID: {problem_id}. Updating notes.")
            notes = extract_actions(payload)
            update_problem_notes(problem_id, notes)
        else:
            logging.info(f"No existing problem found for Event ID: {event_id}. Creating new problem.")
            problem_id = send_to_service_desk(
                title=title,
                description=description,
                priority_name="High",
                status_name="Open",
                impact_details="Impact of detection requires investigation.",
                root_cause="Potential malicious activity detected."
            )
            if problem_id:
                notes = extract_actions(payload)
                update_problem_notes(problem_id, notes)

        return "Processed Successfully", 200
    except Exception as e:
        logging.error(f"Error processing webhook: {e}")
        return str(e), 500

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="ManageEngine ServiceDesk Plus API Utility")
    parser.add_argument("--get-token", action="store_true", help="Retrieve and save tokens")
    args = parser.parse_args()

    client_config = load_client_config()
    CLIENT_ID = client_config["client_id"]
    CLIENT_SECRET = client_config["client_secret"]
    AUTH_CODE = client_config.get("code")
    REDIRECT_URI = "https://localhost"

    if args.get_token:
        if AUTH_CODE:
            get_access_token(CLIENT_ID, CLIENT_SECRET, AUTH_CODE, REDIRECT_URI)
        else:
            logging.error("Authorization code not found in client config.")
    else:
        app.run(host="0.0.0.0", port=5000, debug=True)