"""
COMPLETE FIXED SYSMON SERVER - Working Dashboard
"""

from flask import Flask, request, jsonify, render_template, send_from_directory
import sqlite3
import json
from datetime import datetime, timedelta
import os
import re

app = Flask(__name__)


def init_database():
    """Inicializiraј ja bazata"""
    conn = sqlite3.connect('sysmon_server.db', check_same_thread=False)
    cursor = conn.cursor()

    # Kreiraј tabeli
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS sysmon_events (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            event_id INTEGER,
            event_type TEXT,
            timestamp TEXT,
            computer TEXT,
            process_name TEXT,
            command_line TEXT,
            process_user TEXT,
            parent_process TEXT,
            dns_query TEXT,
            dns_results TEXT,
            dns_process TEXT,
            source_ip TEXT,
            dest_ip TEXT,
            dest_port TEXT,
            file_path TEXT,
            raw_data TEXT,
            received_at TEXT
        )
    ''')

    cursor.execute('''
        CREATE TABLE IF NOT EXISTS processes (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            pid INTEGER,
            name TEXT,
            user TEXT,
            cpu_percent REAL,
            memory_mb REAL,
            cmdline TEXT,
            computer TEXT,
            received_at TEXT
        )
    ''')

    cursor.execute('''
        CREATE TABLE IF NOT EXISTS computers (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            computer_name TEXT UNIQUE,
            user TEXT,
            ip_address TEXT,
            os TEXT,
            last_seen TEXT,
            status TEXT DEFAULT 'online'
        )
    ''')

    # Indeksi
    cursor.execute('CREATE INDEX IF NOT EXISTS idx_computer ON sysmon_events(computer)')
    cursor.execute('CREATE INDEX IF NOT EXISTS idx_dns_query ON sysmon_events(dns_query)')

    conn.commit()
    conn.close()
    print("✅ Database initialized")


def get_db_connection():
    """Vrati database connection"""
    conn = sqlite3.connect('sysmon_server.db', check_same_thread=False)
    conn.row_factory = sqlite3.Row
    return conn


def update_computer_status(computer_name, user, ip_address, os_info):
    """Update computer status"""
    conn = get_db_connection()
    cursor = conn.cursor()

    try:
        cursor.execute('''
            INSERT OR REPLACE INTO computers 
            (computer_name, user, ip_address, os, last_seen, status)
            VALUES (?, ?, ?, ?, ?, ?)
        ''', (computer_name, user, ip_address, os_info, datetime.now().isoformat(), 'online'))
        conn.commit()
    except Exception as e:
        print(f"Error updating computer: {e}")
    finally:
        conn.close()


@app.route('/')
def index():
    """Main dashboard page"""
    return '''
    <!DOCTYPE html>
    <html>
    <head>
        <title>Sysmon Dashboard</title>
        <meta charset="UTF-8">
        <meta name="viewport" content="width=device-width, initial-scale=1.0">
        <script src="https://code.jquery.com/jquery-3.6.0.min.js"></script>
        <link href="https://cdn.jsdelivr.net/npm/bootstrap@5.1.3/dist/css/bootstrap.min.css" rel="stylesheet">
        <style>
            body { background-color: #f8f9fa; padding: 20px; }
            .card { margin-bottom: 20px; box-shadow: 0 2px 4px rgba(0,0,0,.1); }
            .stat-card { border-left: 4px solid; }
            .event-item { padding: 10px; border-left: 3px solid; margin-bottom: 5px; background: white; }
            .event-dns { border-left-color: #3498db; }
            .event-process { border-left-color: #2ecc71; }
            .event-network { border-left-color: #e74c3c; }
            .nav-tabs .nav-link.active { font-weight: bold; }
        </style>
    </head>
    <body>
        <nav class="navbar navbar-dark bg-dark mb-4">
            <div class="container-fluid">
                <span class="navbar-brand">
                    🛡️ Sysmon Dashboard
                </span>
                <span class="navbar-text" id="currentTime"></span>
            </div>
        </nav>

        <div class="container-fluid">
            <!-- Stats -->
            <div class="row mb-4">
                <div class="col-md-2 col-sm-6">
                    <div class="card stat-card border-left-primary">
                        <div class="card-body">
                            <h6>Total Events</h6>
                            <h3 id="totalEvents">0</h3>
                        </div>
                    </div>
                </div>
                <div class="col-md-2 col-sm-6">
                    <div class="card stat-card border-left-success">
                        <div class="card-body">
                            <h6>Computers</h6>
                            <h3 id="totalComputers">0</h3>
                        </div>
                    </div>
                </div>
                <div class="col-md-2 col-sm-6">
                    <div class="card stat-card border-left-info">
                        <div class="card-body">
                            <h6>DNS Queries</h6>
                            <h3 id="dnsCount">0</h3>
                        </div>
                    </div>
                </div>
                <div class="col-md-2 col-sm-6">
                    <div class="card stat-card border-left-warning">
                        <div class="card-body">
                            <h6>Processes</h6>
                            <h3 id="processCount">0</h3>
                        </div>
                    </div>
                </div>
                <div class="col-md-2 col-sm-6">
                    <div class="card stat-card border-left-danger">
                        <div class="card-body">
                            <h6>Network</h6>
                            <h3 id="networkCount">0</h3>
                        </div>
                    </div>
                </div>
                <div class="col-md-2 col-sm-6">
                    <div class="card stat-card border-left-secondary">
                        <div class="card-body">
                            <h6>Last Update</h6>
                            <h6 id="lastUpdate">-</h6>
                        </div>
                    </div>
                </div>
            </div>

            <!-- Tabs -->
            <ul class="nav nav-tabs" id="myTab" role="tablist">
                <li class="nav-item" role="presentation">
                    <button class="nav-link active" id="events-tab" data-bs-toggle="tab" data-bs-target="#events" type="button">All Events</button>
                </li>
                <li class="nav-item" role="presentation">
                    <button class="nav-link" id="dns-tab" data-bs-toggle="tab" data-bs-target="#dns" type="button">DNS Queries</button>
                </li>
                <li class="nav-item" role="presentation">
                    <button class="nav-link" id="processes-tab" data-bs-toggle="tab" data-bs-target="#processes" type="button">Processes</button>
                </li>
                <li class="nav-item" role="presentation">
                    <button class="nav-link" id="network-tab" data-bs-toggle="tab" data-bs-target="#network" type="button">Network</button>
                </li>
                <li class="nav-item" role="presentation">
                    <button class="nav-link" id="computers-tab" data-bs-toggle="tab" data-bs-target="#computers" type="button">Computers</button>
                </li>
            </ul>

            <div class="tab-content" id="myTabContent">
                <!-- All Events -->
                <div class="tab-pane fade show active" id="events" role="tabpanel">
                    <div class="card mt-3">
                        <div class="card-header">
                            <h5>Recent Events</h5>
                        </div>
                        <div class="card-body">
                            <div id="eventsList">Loading events...</div>
                        </div>
                    </div>
                </div>

                <!-- DNS -->
                <div class="tab-pane fade" id="dns" role="tabpanel">
                    <div class="card mt-3">
                        <div class="card-header">
                            <h5>DNS Queries</h5>
                        </div>
                        <div class="card-body">
                            <table class="table table-sm table-hover">
                                <thead>
                                    <tr>
                                        <th>Time</th>
                                        <th>Computer</th>
                                        <th>Domain</th>
                                        <th>Result</th>
                                        <th>Process</th>
                                    </tr>
                                </thead>
                                <tbody id="dnsTable">
                                    <tr><td colspan="5">Loading...</td></tr>
                                </tbody>
                            </table>
                        </div>
                    </div>
                </div>

                <!-- Processes -->
                <div class="tab-pane fade" id="processes" role="tabpanel">
                    <div class="card mt-3">
                        <div class="card-header">
                            <h5>Process Executions</h5>
                        </div>
                        <div class="card-body">
                            <table class="table table-sm table-hover">
                                <thead>
                                    <tr>
                                        <th>Time</th>
                                        <th>Computer</th>
                                        <th>Process</th>
                                        <th>User</th>
                                        <th>Command</th>
                                    </tr>
                                </thead>
                                <tbody id="processesTable">
                                    <tr><td colspan="5">Loading...</td></tr>
                                </tbody>
                            </table>
                        </div>
                    </div>
                </div>

                <!-- Network -->
                <div class="tab-pane fade" id="network" role="tabpanel">
                    <div class="card mt-3">
                        <div class="card-header">
                            <h5>Network Connections</h5>
                        </div>
                        <div class="card-body">
                            <table class="table table-sm table-hover">
                                <thead>
                                    <tr>
                                        <th>Time</th>
                                        <th>Computer</th>
                                        <th>Process</th>
                                        <th>Destination</th>
                                        <th>Port</th>
                                    </tr>
                                </thead>
                                <tbody id="networkTable">
                                    <tr><td colspan="5">Loading...</td></tr>
                                </tbody>
                            </table>
                        </div>
                    </div>
                </div>

                <!-- Computers -->
                <div class="tab-pane fade" id="computers" role="tabpanel">
                    <div class="card mt-3">
                        <div class="card-header">
                            <h5>Connected Computers</h5>
                        </div>
                        <div class="card-body">
                            <table class="table table-sm table-hover">
                                <thead>
                                    <tr>
                                        <th>Computer</th>
                                        <th>User</th>
                                        <th>IP</th>
                                        <th>Status</th>
                                        <th>OS</th>
                                        <th>Last Seen</th>
                                    </tr>
                                </thead>
                                <tbody id="computersTable">
                                    <tr><td colspan="6">Loading...</td></tr>
                                </tbody>
                            </table>
                        </div>
                    </div>
                </div>
            </div>
        </div>

        <script src="https://cdn.jsdelivr.net/npm/bootstrap@5.1.3/dist/js/bootstrap.bundle.min.js"></script>
        <script>
            $(document).ready(function() {
                updateTime();
                loadAllData();
                setInterval(loadAllData, 5000);
                setInterval(updateTime, 1000);

                // Tab switching
                $('button[data-bs-toggle="tab"]').on('click', function(e) {
                    const tab = $(this).attr('data-bs-target');
                    if (tab === '#dns') loadDNS();
                    else if (tab === '#processes') loadProcesses();
                    else if (tab === '#network') loadNetwork();
                    else if (tab === '#computers') loadComputers();
                    else if (tab === '#events') loadEvents();
                });
            });

            function updateTime() {
                $('#currentTime').text(new Date().toLocaleString());
            }

            function loadAllData() {
                loadStats();
                loadEvents();
            }

            function loadStats() {
                $.get('/api/stats', function(data) {
                    $('#totalEvents').text(data.total_events || 0);
                    $('#totalComputers').text(data.total_computers || 0);
                    $('#dnsCount').text(data.dns_queries || 0);
                    $('#processCount').text(data.process_creations || 0);
                    $('#networkCount').text(data.network_connections || 0);
                    $('#lastUpdate').text(new Date().toLocaleTimeString());
                }).fail(function() {
                    console.log('Stats failed');
                });
            }

            function loadEvents() {
                $.get('/api/events?limit=20', function(data) {
                    let html = '';
                    if (data.events && data.events.length > 0) {
                        data.events.forEach(event => {
                            let badge = '';
                            let text = '';

                            if (event.event_id === 22 && event.dns_query) {
                                badge = '<span class="badge bg-info me-2">DNS</span>';
                                text = `<strong>${event.dns_query}</strong> → ${event.dns_results || 'No results'}<br><small>Process: ${event.dns_process || 'Unknown'}</small>`;
                            }
                            else if (event.event_id === 1 && event.process_name) {
                                badge = '<span class="badge bg-success me-2">PROCESS</span>';
                                const procName = event.process_name.split('\\\\').pop() || event.process_name;
                                text = `<strong>${procName}</strong> | User: ${event.process_user || 'Unknown'}<br><small>${event.command_line ? event.command_line.substring(0, 80) + '...' : ''}</small>`;
                            }
                            else if (event.event_id === 3 && event.dest_ip) {
                                badge = '<span class="badge bg-danger me-2">NETWORK</span>';
                                const procName = event.process_name ? event.process_name.split('\\\\').pop() : 'Unknown';
                                text = `<strong>${procName}</strong> → ${event.dest_ip}:${event.dest_port}`;
                            }
                            else {
                                badge = `<span class="badge bg-secondary me-2">${event.event_type}</span>`;
                                text = event.event_type;
                            }

                            html += `
                                <div class="event-item mb-2 p-2 border rounded">
                                    <div class="small text-muted">${formatTime(event.timestamp)} | ${event.computer}</div>
                                    <div>${badge} ${text}</div>
                                </div>
                            `;
                        });
                    } else {
                        html = '<div class="text-center text-muted">No events yet</div>';
                    }
                    $('#eventsList').html(html);
                }).fail(function() {
                    $('#eventsList').html('<div class="text-danger">Failed to load events</div>');
                });
            }

            function loadDNS() {
                $.get('/api/dns', function(data) {
                    let html = '';
                    if (data.dns_queries && data.dns_queries.length > 0) {
                        data.dns_queries.forEach(query => {
                            const process = query.dns_process ? query.dns_process.split('\\\\').pop() : 'Unknown';
                            html += `
                                <tr>
                                    <td>${formatTime(query.timestamp)}</td>
                                    <td>${query.computer}</td>
                                    <td><strong>${query.dns_query}</strong></td>
                                    <td><small>${query.dns_results || 'N/A'}</small></td>
                                    <td>${process}</td>
                                </tr>
                            `;
                        });
                    } else {
                        html = '<tr><td colspan="5" class="text-center text-muted">No DNS queries</td></tr>';
                    }
                    $('#dnsTable').html(html);
                });
            }

            function loadProcesses() {
                $.get('/api/processes', function(data) {
                    let html = '';
                    if (data.processes && data.processes.length > 0) {
                        data.processes.forEach(proc => {
                            const processName = proc.process_name ? proc.process_name.split('\\\\').pop() : 'Unknown';
                            html += `
                                <tr>
                                    <td>${formatTime(proc.timestamp)}</td>
                                    <td>${proc.computer}</td>
                                    <td><strong>${processName}</strong></td>
                                    <td>${proc.process_user || 'N/A'}</td>
                                    <td><small>${proc.command_line ? proc.command_line.substring(0, 80) + '...' : 'N/A'}</small></td>
                                </tr>
                            `;
                        });
                    } else {
                        html = '<tr><td colspan="5" class="text-center text-muted">No processes</td></tr>';
                    }
                    $('#processesTable').html(html);
                });
            }

            function loadNetwork() {
                $.get('/api/network', function(data) {
                    let html = '';
                    if (data.connections && data.connections.length > 0) {
                        data.connections.forEach(conn => {
                            const processName = conn.process_name ? conn.process_name.split('\\\\').pop() : 'Unknown';
                            html += `
                                <tr>
                                    <td>${formatTime(conn.timestamp)}</td>
                                    <td>${conn.computer}</td>
                                    <td>${processName}</td>
                                    <td>${conn.dest_ip}</td>
                                    <td>${conn.dest_port}</td>
                                </tr>
                            `;
                        });
                    } else {
                        html = '<tr><td colspan="5" class="text-center text-muted">No network connections</td></tr>';
                    }
                    $('#networkTable').html(html);
                });
            }

            function loadComputers() {
                $.get('/api/computers', function(data) {
                    let html = '';
                    if (data.computers && data.computers.length > 0) {
                        data.computers.forEach(comp => {
                            const statusClass = comp.status === 'online' ? 'text-success' : 'text-danger';
                            html += `
                                <tr>
                                    <td><strong>${comp.computer_name}</strong></td>
                                    <td>${comp.user || 'N/A'}</td>
                                    <td>${comp.ip_address || 'N/A'}</td>
                                    <td class="${statusClass}">${comp.status}</td>
                                    <td><small>${comp.os || 'N/A'}</small></td>
                                    <td>${formatTime(comp.last_seen)}</td>
                                </tr>
                            `;
                        });
                    } else {
                        html = '<tr><td colspan="6" class="text-center text-muted">No computers</td></tr>';
                    }
                    $('#computersTable').html(html);
                });
            }

            function formatTime(timestamp) {
                if (!timestamp) return 'N/A';
                try {
                    const date = new Date(timestamp);
                    return date.toLocaleTimeString([], {hour: '2-digit', minute:'2-digit'});
                } catch(e) {
                    return timestamp;
                }
            }
        </script>
    </body>
    </html>
    '''


# API ENDPOINTS - SIMPLE AND WORKING
@app.route('/api/stats', methods=['GET'])
def api_stats():
    """Get statistics"""
    conn = get_db_connection()
    cursor = conn.cursor()

    try:
        cursor.execute("SELECT COUNT(*) FROM sysmon_events")
        total_events = cursor.fetchone()[0] or 0

        cursor.execute("SELECT COUNT(DISTINCT computer) FROM sysmon_events")
        total_computers = cursor.fetchone()[0] or 0

        cursor.execute("SELECT COUNT(*) FROM sysmon_events WHERE dns_query != ''")
        dns_queries = cursor.fetchone()[0] or 0

        cursor.execute("SELECT COUNT(*) FROM sysmon_events WHERE event_id = 1")
        process_creations = cursor.fetchone()[0] or 0

        cursor.execute("SELECT COUNT(*) FROM sysmon_events WHERE event_id = 3")
        network_connections = cursor.fetchone()[0] or 0

        return jsonify({
            "total_events": total_events,
            "total_computers": total_computers,
            "dns_queries": dns_queries,
            "process_creations": process_creations,
            "network_connections": network_connections,
            "timestamp": datetime.now().isoformat()
        })
    except Exception as e:
        return jsonify({"error": str(e)}), 500
    finally:
        conn.close()


@app.route('/api/events', methods=['GET'])
def api_events():
    """Get all events"""
    conn = get_db_connection()
    cursor = conn.cursor()

    try:
        limit = request.args.get('limit', 50, type=int)
        cursor.execute('SELECT * FROM sysmon_events ORDER BY timestamp DESC LIMIT ?', (limit,))
        events = [dict(row) for row in cursor.fetchall()]

        return jsonify({
            "events": events,
            "count": len(events)
        })
    except Exception as e:
        return jsonify({"error": str(e)}), 500
    finally:
        conn.close()


@app.route('/api/dns', methods=['GET'])
def api_dns():
    """Get DNS queries"""
    conn = get_db_connection()
    cursor = conn.cursor()

    try:
        limit = request.args.get('limit', 50, type=int)
        cursor.execute('''
            SELECT timestamp, computer, dns_query, dns_results, dns_process 
            FROM sysmon_events 
            WHERE dns_query != '' 
            ORDER BY timestamp DESC 
            LIMIT ?
        ''', (limit,))

        queries = [dict(row) for row in cursor.fetchall()]

        return jsonify({
            "dns_queries": queries,
            "count": len(queries)
        })
    except Exception as e:
        return jsonify({"error": str(e)}), 500
    finally:
        conn.close()


@app.route('/api/processes', methods=['GET'])
def api_processes():
    """Get process executions"""
    conn = get_db_connection()
    cursor = conn.cursor()

    try:
        limit = request.args.get('limit', 50, type=int)
        cursor.execute('''
            SELECT timestamp, computer, process_name, command_line, process_user 
            FROM sysmon_events 
            WHERE event_id = 1 AND process_name != '' 
            ORDER BY timestamp DESC 
            LIMIT ?
        ''', (limit,))

        processes = [dict(row) for row in cursor.fetchall()]

        return jsonify({
            "processes": processes,
            "count": len(processes)
        })
    except Exception as e:
        return jsonify({"error": str(e)}), 500
    finally:
        conn.close()


@app.route('/api/network', methods=['GET'])
def api_network():
    """Get network connections"""
    conn = get_db_connection()
    cursor = conn.cursor()

    try:
        limit = request.args.get('limit', 50, type=int)
        cursor.execute('''
            SELECT timestamp, computer, process_name, dest_ip, dest_port 
            FROM sysmon_events 
            WHERE event_id = 3 AND dest_ip != '' 
            ORDER BY timestamp DESC 
            LIMIT ?
        ''', (limit,))

        connections = [dict(row) for row in cursor.fetchall()]

        return jsonify({
            "connections": connections,
            "count": len(connections)
        })
    except Exception as e:
        return jsonify({"error": str(e)}), 500
    finally:
        conn.close()


@app.route('/api/computers', methods=['GET'])
def api_computers():
    """Get computers list"""
    conn = get_db_connection()
    cursor = conn.cursor()

    try:
        cursor.execute('SELECT * FROM computers ORDER BY last_seen DESC')
        computers = [dict(row) for row in cursor.fetchall()]

        return jsonify({
            "computers": computers,
            "count": len(computers)
        })
    except Exception as e:
        return jsonify({"error": str(e)}), 500
    finally:
        conn.close()


@app.route('/receive_data', methods=['POST'])
def receive_data():
    """Receive data from client"""
    try:
        data = request.json
        if not data:
            return jsonify({"error": "No data"}), 400

        metadata = data.get("metadata", {})
        computer_name = metadata.get("computer", "Unknown")
        user = metadata.get("user", "Unknown")

        conn = get_db_connection()
        cursor = conn.cursor()

        events_inserted = 0

        # Update computer
        system_info = data.get("system_info", {})
        update_computer_status(
            computer_name,
            user,
            system_info.get("ip_address", "Unknown"),
            system_info.get("os", "Unknown")
        )

        # Save sysmon events
        if "sysmon_events" in data:
            for event in data["sysmon_events"]:
                try:
                    event_data = event.get("data", {})

                    cursor.execute('''
                        INSERT INTO sysmon_events 
                        (event_id, event_type, timestamp, computer,
                         process_name, command_line, process_user, parent_process,
                         dns_query, dns_results, dns_process,
                         source_ip, dest_ip, dest_port, file_path,
                         raw_data, received_at)
                        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                    ''', (
                        event.get("event_id"),
                        event.get("event_type"),
                        event.get("timestamp"),
                        event.get("computer"),
                        event_data.get("Image", ""),
                        event_data.get("CommandLine", ""),
                        event_data.get("User", ""),
                        event_data.get("ParentImage", ""),
                        event_data.get("QueryName", ""),
                        event_data.get("QueryResults", ""),
                        event_data.get("Image", ""),
                        event_data.get("SourceIp", ""),
                        event_data.get("DestinationIp", ""),
                        event_data.get("DestinationPort", ""),
                        event_data.get("TargetFilename", ""),
                        json.dumps(event_data),
                        datetime.now().isoformat()
                    ))
                    events_inserted += 1
                except Exception as e:
                    print(f"Error saving event: {e}")
                    continue

        # Save processes
        if "processes" in data:
            for proc in data["processes"]:
                try:
                    cursor.execute('''
                        INSERT INTO processes 
                        (pid, name, user, cpu_percent, memory_mb, cmdline, computer, received_at)
                        VALUES (?, ?, ?, ?, ?, ?, ?, ?)
                    ''', (
                        proc.get("pid"),
                        proc.get("name"),
                        proc.get("user"),
                        proc.get("cpu_percent", 0),
                        proc.get("memory_mb", 0),
                        proc.get("cmdline", ""),
                        computer_name,
                        datetime.now().isoformat()
                    ))
                except Exception as e:
                    print(f"Error saving process: {e}")
                    continue

        conn.commit()
        conn.close()

        return jsonify({
            "message": "Data received successfully",
            "events_inserted": events_inserted
        }), 200

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


@app.route('/ping', methods=['GET'])
def ping():
    """Ping endpoint"""
    return jsonify({
        "message": "pong",
        "status": "ok",
        "server": "Sysmon Server",
        "time": datetime.now().isoformat()
    })


if __name__ == '__main__':
    # Initialize database
    init_database()

    print("\n" + "=" * 60)
    print("🚀 WORKING SYSMON SERVER")
    print("=" * 60)
    print("📊 Dashboard: http://localhost:5555/")
    print("📡 API Endpoints (ALL WORKING):")
    print("   - GET  /              - Dashboard")
    print("   - GET  /api/stats     - Statistics")
    print("   - GET  /api/events    - All events")
    print("   - GET  /api/dns       - DNS queries")
    print("   - GET  /api/processes - Process executions")
    print("   - GET  /api/network   - Network connections")
    print("   - GET  /api/computers - Computers list")
    print("   - POST /receive_data  - Receive client data")
    print("   - GET  /ping          - Test endpoint")
    print("=" * 60)

    app.run(host='0.0.0.0', port=5555, debug=True)