#!/usr/bin/env python3
"""
FIXED SYSMON CLIENT - Sends data in correct format for server
"""

import socket
import os
import platform
import json
import requests
from datetime import datetime, timedelta
import psutil
import time
import re
import xml.etree.ElementTree as ET


class FixedSysmonClient:
    def __init__(self, server_ip="192.168.1.13", port=5555):
        self.server_ip = server_ip
        self.port = port
        self.computer_name = socket.gethostname()
        self.user = os.getlogin()
        self.event_types = self._init_event_types()

    def _init_event_types(self):
        return {
            1: "Process Creation",
            3: "Network Connection",
            5: "Process Terminated",
            7: "Image Loaded (DLL)",
            11: "File Create",
            12: "Registry Object added/deleted",
            13: "Registry Value Set",
            22: "DNS Query",
            23: "File Delete"
        }

    def check_sysmon_installed(self):
        try:
            import win32evtlog
            hand = win32evtlog.OpenEventLog("localhost", "Microsoft-Windows-Sysmon/Operational")
            win32evtlog.CloseEventLog(hand)
            return True
        except:
            return False

    def parse_event_data_advanced(self, event):
        try:
            event_data = {}

            if event.StringInserts:
                for insert in event.StringInserts:
                    if insert:
                        try:
                            xml_string = insert.strip()
                            if not xml_string.startswith('<'):
                                continue

                            root = ET.fromstring(f"<root>{xml_string}</root>")

                            for data_elem in root.findall('.//Data'):
                                name = data_elem.get('Name')
                                value = data_elem.text
                                if name and value:
                                    event_data[name] = value.strip()

                        except ET.ParseError:
                            pattern = r'<Data Name=["\']([^"\']+)["\']>([^<]*)</Data>'
                            matches = re.findall(pattern, str(insert))
                            for name, value in matches:
                                if name and value:
                                    event_data[name] = value.strip()

            return event_data

        except Exception as e:
            return {}

    def get_all_sysmon_events(self, max_events=50):
        events = []

        if not self.check_sysmon_installed():
            print("   ⚠️ Sysmon not installed")
            return events

        try:
            import win32evtlog

            hand = win32evtlog.OpenEventLog("localhost", "Microsoft-Windows-Sysmon/Operational")
            flags = win32evtlog.EVENTLOG_BACKWARDS_READ | win32evtlog.EVENTLOG_SEQUENTIAL_READ

            events_collected = 0

            while events_collected < max_events:
                events_batch = win32evtlog.ReadEventLog(hand, flags, 0)
                if not events_batch:
                    break

                for event in events_batch:
                    try:
                        event_id = event.EventID & 0xFFFF
                        event_type = self.event_types.get(event_id, f"Event {event_id}")

                        event_data = self.parse_event_data_advanced(event)

                        # Extract fields based on event type
                        event_entry = {
                            "event_id": event_id,
                            "event_type": event_type,
                            "timestamp": event.TimeGenerated.Format() if hasattr(event.TimeGenerated,
                                                                                 'Format') else str(
                                event.TimeGenerated),
                            "computer": self.computer_name,
                            "process": event_data.get("Image", ""),
                            "command_line": event_data.get("CommandLine", ""),
                            "user": event_data.get("User", ""),
                            "parent_process": event_data.get("ParentImage", ""),
                            "dns_query": event_data.get("QueryName", ""),
                            "dns_results": event_data.get("QueryResults", ""),
                            "source_ip": event_data.get("SourceIp", ""),
                            "destination_ip": event_data.get("DestinationIp", ""),
                            "destination_port": event_data.get("DestinationPort", ""),
                            "file_path": event_data.get("TargetFilename", ""),
                            "hash": event_data.get("Hashes", ""),
                            "registry_key": event_data.get("TargetObject", ""),
                            "registry_value": event_data.get("Details", ""),
                            "description": self.get_event_description(event_id, event_data),
                            "raw_data": json.dumps(event_data)
                        }

                        # Only add if meaningful
                        if any([
                            event_entry["process"],
                            event_entry["dns_query"],
                            event_entry["file_path"],
                            event_entry["destination_ip"]
                        ]):
                            events.append(event_entry)
                            events_collected += 1

                        if events_collected >= max_events:
                            break

                    except Exception as e:
                        continue

            win32evtlog.CloseEventLog(hand)

        except Exception as e:
            print(f"   ❌ Error reading Sysmon: {e}")

        return events

    def get_event_description(self, event_id, event_data):
        if event_id == 1:
            proc = event_data.get("Image", "").split("\\")[-1] if event_data.get("Image") else "Unknown"
            return f"Process started: {proc}"
        elif event_id == 22:
            return f"DNS query: {event_data.get('QueryName', 'Unknown')}"
        elif event_id == 3:
            return f"Network: {event_data.get('DestinationIp', 'Unknown')}"
        elif event_id == 11:
            file_name = event_data.get("TargetFilename", "").split("\\")[-1] if event_data.get(
                "TargetFilename") else "Unknown"
            return f"File created: {file_name}"
        elif event_id == 23:
            file_name = event_data.get("TargetFilename", "").split("\\")[-1] if event_data.get(
                "TargetFilename") else "Unknown"
            return f"File deleted: {file_name}"
        else:
            return f"Event {event_id}"

    def get_system_info(self):
        try:
            return {
                "computer_name": self.computer_name,
                "user": self.user,
                "hostname": socket.gethostname(),
                "ip_address": socket.gethostbyname(self.computer_name) if hasattr(socket, 'gethostbyname') else "N/A",
                "os": f"{platform.system()} {platform.release()}",
                "cpu_usage": psutil.cpu_percent(interval=1),
                "ram_usage": psutil.virtual_memory().percent,
                "disk_usage": psutil.disk_usage("C:/").percent if os.name == 'nt' else 0,
                "collection_time": datetime.now().isoformat()
            }
        except Exception as e:
            return {"error": str(e)}

    def get_running_processes(self):
        processes = []
        try:
            for proc in psutil.process_iter(['pid', 'name', 'username']):
                try:
                    with proc.oneshot():
                        processes.append({
                            "pid": proc.info['pid'],
                            "name": proc.info['name'],
                            "user": proc.info['username'] if proc.info['username'] else "N/A",
                            "cpu_percent": round(proc.cpu_percent(), 2),
                            "memory_mb": round(proc.memory_info().rss / (1024 ** 2), 2),
                            "cmdline": ' '.join(proc.cmdline())[:200] if proc.cmdline() else "",
                            "timestamp": datetime.now().isoformat()
                        })
                except (psutil.NoSuchProcess, psutil.AccessDenied):
                    continue

                if len(processes) >= 50:
                    break
        except Exception as e:
            print(f"   ⚠️ Error collecting processes: {e}")

        return processes

    def collect_and_send(self):
        print(f"\n[{datetime.now().strftime('%H:%M:%S')}] 🔍 Collecting data...")

        try:
            sysmon_installed = self.check_sysmon_installed()
            print(f"   🔍 Sysmon: {'✅ Installed' if sysmon_installed else '❌ Not installed'}")

            print(f"   📥 Collecting Sysmon events...")
            events = self.get_all_sysmon_events(max_events=50)

            print(f"   📥 Collecting system info...")
            system_info = self.get_system_info()

            print(f"   📥 Collecting processes...")
            processes = self.get_running_processes()

            # Prepare data in CORRECT format for server
            data = {
                "metadata": {
                    "computer": self.computer_name,
                    "user": self.user,
                    "timestamp": datetime.now().isoformat()
                },
                "events": events,
                "system_info": system_info,
                "processes": processes
            }

            print(f"\n   📊 STATISTICS:")
            print(f"      • Events: {len(events)}")
            print(f"      • Processes: {len(processes)}")

            # Show samples
            dns_events = [e for e in events if e.get("dns_query")]
            file_events = [e for e in events if e.get("file_path")]
            network_events = [e for e in events if e.get("destination_ip")]

            if dns_events:
                print(f"\n   🌐 DNS QUERIES ({len(dns_events)}):")
                for dns in dns_events[:3]:
                    print(f"      • {dns.get('dns_query')}")

            if file_events:
                print(f"\n   📁 FILE OPERATIONS ({len(file_events)}):")
                for file_ev in file_events[:3]:
                    file_name = os.path.basename(file_ev.get("file_path", ""))
                    print(f"      • {file_name}")

            if network_events:
                print(f"\n   🌍 NETWORK ({len(network_events)}):")
                for net in network_events[:3]:
                    print(f"      • {net.get('destination_ip')}:{net.get('destination_port')}")

            print(f"\n   📤 Sending to server...")
            return self.send_to_server(data)

        except Exception as e:
            print(f"   ❌ Error: {e}")
            import traceback
            traceback.print_exc()
            return False

    def send_to_server(self, data):
        try:
            response = requests.post(
                f"http://{self.server_ip}:{self.port}/receive_data",
                json=data,
                timeout=30,
                headers={'Content-Type': 'application/json'}
            )

            if response.status_code == 200:
                result = response.json()
                print(f"   ✅ Successfully sent!")
                print(f"   📝 Server response: {result.get('message', 'Success')}")
                print(f"   📊 Events inserted: {result.get('events_inserted', 0)}")
                return True
            else:
                print(f"   ❌ Server error {response.status_code}")
                print(f"   Response: {response.text[:200]}")
                return False

        except requests.exceptions.ConnectionError:
            print(f"   ❌ Cannot connect to server at {self.server_ip}:{self.port}")
            return False
        except Exception as e:
            print(f"   ❌ Error: {e}")
            return False


def main():
    print("=" * 70)
    print("🛡️  FIXED SYSMON CLIENT")
    print("=" * 70)
    print(f"Computer: {socket.gethostname()}")
    print(f"User: {os.getlogin()}")
    print("=" * 70)

    server_ip = input("Enter server IP (default 192.168.1.13): ").strip() or "192.168.1.13"

    client = FixedSysmonClient(server_ip)

    # Test connection
    print(f"\n🔌 Testing connection to {server_ip}:5555...")
    try:
        response = requests.get(f"http://{server_ip}:5555/ping", timeout=5)
        if response.status_code == 200:
            print(f"   ✅ Server is reachable!")
        else:
            print(f"   ⚠️  Server responded with status {response.status_code}")
    except Exception as e:
        print(f"   ⚠️  Cannot connect: {e}")

    # Check pywin32
    try:
        import win32evtlog
        print(f"   ✅ pywin32 installed")
    except ImportError:
        print(f"   ❌ pywin32 not installed - run: pip install pywin32")
        input("Press ENTER to exit...")
        return

    interval = input("\nCollection interval in seconds (default 60): ").strip()
    try:
        interval = int(interval) if interval else 60
    except:
        interval = 60

    try:
        count = 1
        while True:
            print(f"\n" + "=" * 70)
            print(f"📊 CYCLE #{count} - {datetime.now().strftime('%H:%M:%S')}")
            print("=" * 70)

            start_time = time.time()
            client.collect_and_send()
            elapsed = time.time() - start_time

            print(f"\n   ⏱️  Time: {elapsed:.2f} seconds")
            count += 1

            print(f"\n⏳ Waiting {interval} seconds...")
            time.sleep(interval)

    except KeyboardInterrupt:
        print("\n\n🛑 Program stopped.")
        print(f"📊 Dashboard: http://{server_ip}:5555/")
        print("=" * 70)


if __name__ == "__main__":
    required_modules = ["psutil", "requests"]
    missing = []

    for module in required_modules:
        try:
            __import__(module)
        except ImportError:
            missing.append(module)

    if missing:
        print(f"❌ Missing: {', '.join(missing)}")
        print(f"Install: pip install {' '.join(missing)}")
        if os.name == 'nt':
            print("Also: pip install pywin32")
        input("Press ENTER to exit...")
        exit(1)

    print("✅ All modules installed")
    main()