from flask import Flask, request, jsonify
import requests
import json
from datetime import datetime

app = Flask(__name__)

@app.route('/webhook', methods=['POST'])
def webhook():
    try:
        payload = request.json
        zt_event = {
            "event": {
                "user": payload.get("user"),
                "state": payload.get("state"),
                "userName": payload.get("userName"),
                "protocol": payload.get("protocol"),
                "eventSource": payload.get("eventSource"),
                "numFiles": payload.get("numFiles"),
                "nes": payload.get("nes"),
                "detected": payload.get("detected"),
                "clientIPs": payload.get("clientIPs"),
                "shares": payload.get("shares"),
                "files": payload.get("files"),
            }
        }
        print(json.dumps(zt_event, indent=4))
        
        zt_summary_event = {
            "event": {
                "state": payload.get("state"),
                "userName": payload.get("userName"),
                "shares": payload.get("shares"),
            }
        }
        shares_array = payload.get("shares", [])
        shares_string = ", ".join([item.get('name', '') for item in shares_array])
        zt_summary_event["event"]["shares"] = shares_string
        print(json.dumps(zt_summary_event, indent=4))

        print("Starting inventory retrieval of agents")
        ########################## Trend Vision One Integration section ####################################
        global host_isolation
        host_isolation = 'true'
        global token
        token = 'xxxxxxx'
        url_base = 'https://api.xdr.trendmicro.com/v3.0/endpointSecurity/endpoints'
        
        query_params = {
            'orderBy': 'agentGuid',
            'top': '1000',
            'select': 'select=endpointName,agentGuid,ipAddresses,osName,displayName'
        }
        headers = {
            'Authorization': f'Bearer {token}',
        }
        endpointinventoryjson = requests.get(url_base, params=query_params, headers=headers)
        print(endpointinventoryjson.status_code)
        for k, v in endpointinventoryjson.headers.items():
            print(f'{k}: {v}')
        print('')
        if 'application/json' in endpointinventoryjson.headers.get('Content-Type', '') and len(endpointinventoryjson.content):
            print(json.dumps(endpointinventoryjson.json(), indent=4))
        else:
            print(endpointinventoryjson.text)

        host_client_ips = zt_event["event"]["clientIPs"]
        print("Zero Trust Client IP of infected host returned to lookup in Endpoint inventory:", host_client_ips)

        for item in endpointinventoryjson.json()["items"]:
            for ip in item["ipAddresses"]:
                if ip in host_client_ips:
                    print(f"Found IP Address {ip} in item with agentGuid to be isolated {item['agentGuid']}")
                    agentGUID = item['agentGuid']

        url_base = 'https://api.xdr.trendmicro.com'
        url_path = '/v3.0/response/endpoints/isolate'
        url_API = url_base + url_path
        Isolation_description = 'Superna Zero Trust host isolation request due to Ransomware attack'

        query_params = {}
        headers = {
            'Authorization': f'Bearer {token}',
            'Content-Type': 'application/json;charset=utf-8'
        }
        body = [
            {
                'description': Isolation_description,
                'agentGuid': agentGUID
            }
        ]

        tvoapirequest = requests.post(url_API, params=query_params, headers=headers, json=body)
        print(tvoapirequest.status_code)
        for k, v in tvoapirequest.headers.items():
            print(f'{k}: {v}')
        print('')
        if 'application/json' in tvoapirequest.headers.get('Content-Type', '') and len(tvoapirequest.content):
            print(json.dumps(tvoapirequest.json(), indent=4))
        else:
            print(tvoapirequest.text)

    except Exception as e:
        print(f"Error: {str(e)}")
        return jsonify({"error": str(e)}), 500

    return jsonify({"message": "Webhook processed"}), 200

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