import re
import subprocess
import shlex
import time
import os
import signal
from typing import Dict, Any, List
from src.utils.logger import get_logger

logger = get_logger("nmap_scanner")

class NmapScanner:
    def __init__(self, ip: str, scan_type: str = "Quick Scan", custom_args: str = ""):
        """
        Инициализация сканера Nmap

        :param ip: целевой IP-адрес
        :param scan_type: тип сканирования (Quick Scan, Intensive Scan, Full Scan, Custom Scan)
        :param custom_args: пользовательские аргументы для сканирования
        """
        self.ip = ip
        self.scan_type = scan_type
        self.custom_args = custom_args
        self.process = None
        self.start_time = 0
        self.output_file = "/home/lodqa/attack_module_data/nmap_output.xml"
        self.stats = {
            'hosts_up': 0,
            'hosts_down': 0,
            'duration': 0.0,
            'start_time': 0,
            'end_time': 0,
            'hosts': {},
            'service_info': {}
        }
        self.running = False
        logger.info(f"Инициализация NmapScanner для IP: {ip}, тип: {scan_type}, аргументы: {custom_args}")

    def _get_nmap_args(self) -> List[str]:
        """
        Определение аргументов Nmap на основе типа сканирования и пользовательских аргументов

        :return: список аргументов Nmap
        """
        scan_args = {
            "Quick Scan": "-T4 -F",
            "Intensive Scan": "-T4 -A",
            "Full Scan": "-T4 -p- -A",
            "Custom Scan": self.custom_args
        }
        base_args = scan_args.get(self.scan_type, "-T4")
        args = shlex.split(base_args)
        args.extend(["-oX", self.output_file])
        args.append(self.ip)
        return ["nmap"] + args

    def start_scan(self) -> Dict[str, Any]:
        """
        Запуск сканирования Nmap и возврат результатов

        :return: словарь с результатами сканирования
        """
        logger.info(f"🚀 Запуск сканирования Nmap для {self.ip}")
        self.running = True
        self.start_time = time.time()
        self.stats['start_time'] = self.start_time

        command = self._get_nmap_args()
        logger.info(f"Команда Nmap: {' '.join(command)}")

        try:
            # Запуск nmap в фоновом процессе
            self.process = subprocess.Popen(
                command,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True
            )

            # Ожидание завершения процесса
            while self.running and self.process.poll() is None:
                time.sleep(0.5)

            if self.running:
                # Процесс завершился нормально
                stdout, stderr = self.process.communicate()
                self.stats['end_time'] = time.time()
                self.stats['duration'] = self.stats['end_time'] - self.start_time

                # Чтение вывода из файла
                if os.path.exists(self.output_file):
                    with open(self.output_file, 'r') as f:
                        xml_output = f.read()
                    os.remove(self.output_file)
                    self._parse_output(xml_output, stderr)
                else:
                    self.stats['error'] = "Output file not generated"
                    logger.error("No output file generated by nmap")
            else:
                # Сканирование было остановлено
                self.stats['error'] = "Scan stopped by user"

            self._log_summary()
            return self.stats

        except Exception as e:
            logger.error(f"Ошибка при выполнении Nmap: {e}")
            self.stats['error'] = str(e)
            return self.stats
        finally:
            self.running = False
            self.process = None

    def stop(self):
        """Остановка сканирования Nmap"""
        if self.process and self.running:
            logger.info("🛑 Stopping Nmap scan")
            self.running = False
            self.process.terminate()
            try:
                self.process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                self.process.kill()

    def _parse_output(self, xml_output: str, stderr: str):
        """
        Анализ XML-вывода Nmap

        :param xml_output: XML-вывод команды nmap
        :param stderr: ошибки, если есть
        """
        if stderr:
            logger.error(f"Nmap errors:\n{stderr}")

        if not xml_output:
            logger.warning("No XML output to parse.")
            return

        try:
            import xml.etree.ElementTree as ET
            root = ET.fromstring(xml_output)

            # Подсчет хостов
            for host in root.findall(".//host"):
                status = host.find("status").get("state")
                if status == "up":
                    self.stats['hosts_up'] += 1
                else:
                    self.stats['hosts_down'] += 1

                # Извлечение информации о хосте
                ip = host.find("address[@addrtype='ipv4']").get("addr")
                host_data = {
                    "ip": ip,
                    "state": status,
                    "ports": [],
                    "osmatch": [],
                    "macaddress": {}
                }

                # Извлечение MAC-адреса
                mac = host.find("address[@addrtype='mac']")
                if mac is not None:
                    host_data["macaddress"] = {"addr": mac.get("addr"), "vendor": mac.get("vendor", "Unknown")}

                # Извлечение информации о портах
                ports = host.find("ports")
                if ports is not None:
                    for port in ports.findall("port"):
                        port_info = {
                            "portid": port.get("portid"),
                            "protocol": port.get("protocol"),
                            "state": port.find("state").get("state"),
                            "service": {}
                        }
                        service = port.find("service")
                        if service is not None:
                            port_info["service"] = {
                                "name": service.get("name", "Unknown"),
                                "product": service.get("product", ""),
                                "version": service.get("version", ""),
                                "extrainfo": service.get("extrainfo", "")
                            }
                        host_data["ports"].append(port_info)

                # Извлечение информации об ОС
                os = host.find("os")
                if os is not None:
                    for osmatch in os.findall("osmatch"):
                        host_data["osmatch"].append({
                            "name": osmatch.get("name", "Unknown"),
                            "accuracy": osmatch.get("accuracy", "0")
                        })

                self.stats['hosts'][ip] = host_data

            logger.debug(f"Parsed hosts: {self.stats['hosts']}")

        except Exception as e:
            logger.error(f"Ошибка при разборе XML вывода Nmap: {e}")
            self.stats['error'] = str(e)

    def _log_summary(self):
        """Логирование результатов сканирования"""
        if 'error' in self.stats:
            logger.error(f"❌ Scan failed: {self.stats['error']}")
            return
            
        logger.success("📊 Nmap scan summary:")
        logger.success(f"  Target: {self.ip}")
        logger.success(f"  Scan type: {self.scan_type}")
        logger.success(f"  Hosts up: {self.stats['hosts_up']}")
        logger.success(f"  Hosts down: {self.stats['hosts_down']}")
        logger.success(f"  Duration: {self.stats['duration']:.2f} seconds")