gonets/SRV/services/subnet_service.py

from scapy.config import conf
from mac_vendor_lookup import MacLookup
from typing import List, Optional, Tuple, Dict
import socket
import logging
import netifaces
import ipaddress
from concurrent.futures import ThreadPoolExecutor
import time
from functools import lru_cache
import json
from pathlib import Path
from datetime import datetime, timedelta
import subprocess
from utils.i18n import get_translations, _
from babel.support import Translations
from fastapi import Depends

from models.subnet import SubnetScanResult, Device
from config.ports import port_config
from scapy.layers.inet import IP, TCP
from scapy.layers.l2 import ARP, Ether
from scapy.sendrecv import srp, sr1

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)

class NetworkScanner:
    """Класс для сканирования подсети с использованием ARP-запросов.

    Выполняет обнаружение устройств в указанной подсети через ARP-пакеты.
    """

    def __init__(self, interface: str, translations: Translations = Depends(get_translations("services"))):
        """Инициализирует сканер сети с указанным сетевым интерфейсом.

        Args:
            interface (str): Название сетевого интерфейса (например, "eth0").
        """
        self.interface = interface
        conf.iface = interface
        self.translations = translations

    def scan(self, subnet: str) -> List[Tuple[str, str]]:
        """Сканирует подсеть с использованием ARP-запросов.

        Args:
            subnet (str): Подсеть в формате CIDR (например, "192.168.1.0/24").

        Returns:
            List[Tuple[str, str]]: Список кортежей (IP-адрес, MAC-адрес) обнаруженных
                устройств.

        Raises:
            Exception: Если произошла ошибка при отправке ARP-пакетов.
        """
        try:
            arp = ARP(pdst=subnet)
            ether = Ether(dst="ff:ff:ff:ff:ff:ff")
            packet = ether / arp
            result = srp(packet,
                         timeout=2,
                         verbose=False,
                         iface=self.interface)[0]
            return [(received.psrc, received.hwsrc)
                    for sent, received in result]
        except Exception as e:
            logger.error(f"Error scanning subnet {subnet}: {str(e)}")
            return []

class PortScanner:
    """Класс для сканирования открытых портов на устройствах.

    Выполняет сканирование портов с использованием TCP SYN-запросов.
    """

    def __init__(self, interface: str, translations: Translations = Depends(get_translations("services"))):
        """Инициализирует сканер портов с указанным сетевым интерфейсом.

        Args:
            interface (str): Название сетевого интерфейса (например, "eth0").
        """
        self.interface = interface
        self.translations = translations

    def scan(self, ip: str) -> List[str]:
        """Сканирует открытые порты на указанном IP-адресе.

        Args:
            ip (str): IP-адрес устройства для сканирования.

        Returns:
            List[str]: Список имён открытых портов (например, "HTTP", "SSH").

        Raises:
            Exception: Если произошла ошибка при сканировании портов.
        """
        open_ports = []

        def check_port(port: int) -> Optional[str]:
            """Проверяет доступность конкретного порта.

            Args:
                port (int): Номер порта для проверки.

            Returns:
                Optional[str]: Имя порта, если открыт, иначе None.
            """
            try:
                pkt = sr1(
                    IP(dst=ip) / TCP(dport=port, flags="S"),
                    timeout=2,
                    verbose=False,
                )
                if pkt and pkt.haslayer(TCP) and pkt[TCP].flags == 0x12:
                    return port_config.get_port_name(port)
                return None
            except Exception as e:
                logger.debug(
                    f"Error checking port {port} for IP {ip}: {str(e)}")
                return None

        try:
            max_workers = min(len(port_config.get_ports()), 15)
            with ThreadPoolExecutor(max_workers=max_workers) as executor:
                results = executor.map(check_port, port_config.get_ports())
                open_ports = [port for port in results if port is not None]
            logger.debug(f"Open ports for IP {ip}: {open_ports}")
            return open_ports
        except Exception as e:
            logger.warning(
                f"Error scanning ports for IP {ip}: {str(e)}")
            return []

class DeviceInfoCollector:
    """Класс для сбора информации об устройствах, включая вендоров и хостнеймы.

    Управляет кэшированием данных NetBIOS и DNS.
    """

    def __init__(self, translations: Translations = Depends(get_translations("services"))):
        """Инициализирует коллектор информации об устройствах.

        Создаёт объект MacLookup для определения вендоров и инициализирует кэши.
        """
        self.mac_lookup = MacLookup()
        self.nbtscan_cache = {}
        self.nbtscan_results = {}
        self.translations = translations

    @lru_cache(maxsize=1000)
    def get_vendor(self, mac: str) -> str:
        """Определяет вендора по MAC-адресу с использованием кэширования.

        Args:
            mac (str): MAC-адрес устройства в формате XX:XX:XX:XX:XX:XX.

        Returns:
            str: Название вендора или "Unknown", если вендор не определён.

        Raises:
            Exception: Если произошла ошибка при запросе к базе вендоров.
        """
        try:
            return self.mac_lookup.lookup(mac)
        except Exception as e:
            logger.warning(
                f"Failed to determine vendor for MAC {mac}: {str(e)}")
            return self.translations.gettext("Неизвестный")

    def scan_nbtscan(self, subnet: str) -> None:
        """Выполняет сканирование NetBIOS для указанной подсети.

        Args:
            subnet (str): Подсеть в формате CIDR (например, "192.168.1.0/24").

        Сохраняет результаты в кэш, игнорируя хостнеймы "<unknown>".
        """
        if subnet not in self.nbtscan_results:
            try:
                logger.debug(f"Starting nbtscan for subnet {subnet}")
                result = subprocess.run(["sudo", "nbtscan", "-r", str(subnet)],
                                       capture_output=True,
                                       text=True,
                                       check=False,
                                       timeout=10)
                logger.debug(f"Full nbtscan output: {result.stdout}")
                if result.returncode == 0:
                    for line in result.stdout.splitlines():
                        logger.debug(f"Processing line: {line}")
                        if not line.startswith(("IP address", "-", "Doing NBT")):
                            parts = [p for p in line.split() if p]
                            if len(parts) >= 2:
                                ip_addr, hostname = parts[0], parts[1]
                                if hostname != "<unknown>":
                                    self.nbtscan_cache[ip_addr] = hostname
                                    logger.info(f"NetBIOS: Found hostname {hostname} for IP {ip_addr}")
                self.nbtscan_results[subnet] = True
            except subprocess.TimeoutExpired:
                logger.error(f"NetBIOS: Timeout for subnet {subnet}")
            except Exception as e:
                logger.error(f"NetBIOS: Error for subnet {subnet}: {str(e)}")

    def get_hostname(self, ip: str) -> Optional[str]:
        """Определяет хостнейм для указанного IP-адреса.

        Args:
            ip (str): IP-адрес устройства.

        Returns:
            Optional[str]: Хостнейм, если найден, иначе None.

        Сначала пытается использовать DNS, затем NetBIOS через nbtscan.
        """
        try:
            socket.setdefaulttimeout(2)
            return socket.gethostbyaddr(ip)[0]
        except (socket.herror, socket.gaierror, socket.timeout):
            logger.debug(f"DNS: Failed to resolve name for IP {ip}")

        subnet = ipaddress.ip_network(f"{ip}/24", strict=False)
        if subnet not in self.nbtscan_results:
            self.scan_nbtscan(str(subnet))

        if ip in self.nbtscan_cache:
            return self.nbtscan_cache[ip]

        logger.debug(f"Failed to determine hostname for IP {ip} via DNS or NetBIOS")
        return None

class DeviceStorage:
    """Класс для хранения и управления данными об устройствах.

    Хранит устройства в JSON-файле и удаляет устаревшие записи.
    """

    def __init__(self, storage_path: str = "storage/devices/devices.json", translations: Translations = Depends(get_translations("services"))):
        """Инициализирует хранилище устройств с указанным путём.

        Args:
            storage_path (str, optional): Путь к файлу хранения.
                По умолчанию "storage/devices/devices.json".

        Создаёт родительские директории, если они не существуют.
        """
        self.storage_file = Path(storage_path)
        self.storage_file.parent.mkdir(parents=True, exist_ok=True)
        self.offline_timeout = timedelta(days=7)
        self.translations = translations

    def load_devices(self) -> Dict[str, Device]:
        """Загружает устройства из файла хранения.

        Returns:
            Dict[str, Device]: Словарь устройств, где ключ — MAC-адрес.

        Удаляет устройства с неверным форматом last_seen и возвращает пустой словарь
        при ошибке.
        """
        try:
            if self.storage_file.exists():
                with open(self.storage_file, "r", encoding="utf-8") as f:
                    data = json.load(f)
                    devices = {}
                    for device_data in data.get("devices", []):
                        device = Device(**device_data)
                        if device.mac:
                            devices[device.mac] = device
                    return self._clean_old_devices(devices)
            return {}
        except Exception as e:
            logger.error(f"Error loading devices: {str(e)}")
            return {}

    def _clean_old_devices(self, devices: Dict[str, Device]) -> Dict[str, Device]:
        """Очищает устаревшие устройства из словаря.

        Args:
            devices (Dict[str, Device]): Словарь устройств с MAC-адресами как ключами.

        Returns:
            Dict[str, Device]: Отфильтрованный словарь, исключающий устройства
                offline более 7 дней.
        """
        now = datetime.now()
        filtered = {}
        for mac, device in devices.items():
            if device.status == "offline" and device.last_seen:
                try:
                    last_seen = datetime.fromisoformat(device.last_seen)
                    if now - last_seen <= self.offline_timeout:
                        filtered[mac] = device
                    else:
                        logger.info(f"Removed device {mac} (offline > 7 days)")
                except ValueError:
                    logger.warning(f"Invalid last_seen format for device {mac}")
            else:
                filtered[mac] = device
        return filtered

    def save_devices(self, devices: List[Device]) -> None:
        """Сохраняет список устройств в файл.

        Args:
            devices (List[Device]): Список объектов Device для сохранения.

        Raises:
            Exception: Если произошла ошибка при записи в файл.
        """
        try:
            data = SubnetScanResult(devices=devices).model_dump()
            with open(self.storage_file, "w", encoding="utf-8") as f:
                json.dump(data, f, indent=2, ensure_ascii=False)
            logger.info(f"Saved {len(devices)} devices to {self.storage_file}")
        except Exception as e:
            logger.error(f"Error saving devices: {str(e)}")
            raise Exception(self.translations.gettext("Ошибка при сохранении устройств: {error}").format(error=str(e)))

class SubnetService:
    """Сервис для сканирования подсетей и управления устройствами.

    Объединяет функционал сканирования сети, портов и хранения данных об устройствах.
    """

    def __init__(self, translations: Translations = Depends(get_translations("services"))):
        """Инициализирует сервис подсети с настройкой всех компонентов.

        Автоматически определяет интерфейс и инициализирует сканеры и хранилище.
        """
        self.interface = self._get_default_interface()
        self.network_scanner = NetworkScanner(self.interface)
        self.port_scanner = PortScanner(self.interface)
        self.device_info = DeviceInfoCollector()
        self.storage = DeviceStorage()
        self.translations = translations

    def _get_default_interface(self) -> str:
        """Определяет имя сетевого интерфейса по умолчанию.

        Returns:
            str: Название интерфейса (например, "eth0").

        Использует информацию о шлюзе или первый доступный интерфейс, отличный от "lo".
        """
        try:
            gateways = netifaces.gateways()
            default_gateway = gateways.get("default", {}).get(netifaces.AF_INET)
            if default_gateway:
                return default_gateway[1]
            interfaces = netifaces.interfaces()
            for iface in interfaces:
                if iface != "lo":
                    return iface
            return "eth0"
        except Exception as e:
            logger.warning(f"Failed to determine interface: {str(e)}")
            return "eth0"

    def scan_subnet(self, subnet: str, ports: List[int] = None) -> SubnetScanResult:
        """Сканирует указанную подсеть и возвращает результаты.

        Args:
            subnet (str): Подсеть в формате CIDR (например, "192.168.1.0/24").
            ports (List[int], optional): Список портов для сканирования.
                По умолчанию None (используются порты из config.ports).

        Returns:
            SubnetScanResult: Объект с результатами сканирования, содержащий список
                устройств.

        Сохраняет данные в хранилище и обновляет статус устройств.
        """
        start_time = time.time()
        devices = []
        now = datetime.now().astimezone()
        formatted_now = now.strftime("%d-%m-%Y, %H:%M")
        try:
            saved_devices = self.storage.load_devices()

            local_device = self._get_local_device(subnet)
            if local_device:
                local_device.status = "online"
                local_device.last_seen = formatted_now
                devices.append(local_device)
                if local_device.mac:
                    saved_devices[local_device.mac] = local_device

            discovered_devices = self.network_scanner.scan(subnet)
            self.device_info.scan_nbtscan(subnet)

            def process_device(ip_mac: Tuple[str, str]) -> Optional[Device]:
                """Обрабатывает устройство из пары IP и MAC.

                Args:
                    ip_mac (Tuple[str, str]): Кортеж с IP-адресом и MAC-адресом.

                Returns:
                    Optional[Device]: Объект Device или None, если устройство локальное.
                """
                ip, mac = ip_mac
                if local_device and ip == local_device.ip:
                    return None

                existing_device = saved_devices.get(mac)
                if existing_device:
                    return Device(
                        ip=ip,
                        mac=mac,
                        hostname=self.device_info.get_hostname(ip) or existing_device.hostname,
                        vendor=self.device_info.get_vendor(mac),
                        open_ports=self.port_scanner.scan(ip) or existing_device.open_ports,
                        status="online",
                        last_seen=formatted_now,
                        note=existing_device.note,
                    )
                vendor = self.device_info.get_vendor(mac)
                hostname = self.device_info.get_hostname(ip)
                open_ports = self.port_scanner.scan(ip)
                return Device(
                    ip=ip,
                    mac=mac,
                    hostname=hostname,
                    vendor=vendor,
                    open_ports=open_ports,
                    status="online",
                    last_seen=formatted_now,
                    note=None,
                )

            with ThreadPoolExecutor(max_workers=10) as executor:
                device_results = executor.map(process_device, discovered_devices)
                devices.extend([d for d in device_results if d is not None])

            active_macs = {d.mac for d in devices if d.mac}
            for mac, device in saved_devices.items():
                if mac not in active_macs:
                    device.status = "offline"
                    devices.append(device)

            devices.sort(key=lambda x: int(ipaddress.IPv4Address(x.ip)))

            for device in devices:
                if device.mac:
                    saved_devices[device.mac] = device
            self.storage.save_devices(list(saved_devices.values()))

            logger.info(
                f"Subnet scan {subnet} completed in {time.time() - start_time:.2f} seconds, "
                f"found {len(devices)} devices (active: {len(active_macs)})"
            )
            return SubnetScanResult(devices=devices)

        except Exception as e:
            logger.error(f"Error scanning subnet {subnet}: {str(e)}")
            return SubnetScanResult(devices=devices)

    def _get_local_device(self, subnet: str) -> Optional[Device]:
        """Определяет локальное устройство в указанной подсети.

        Args:
            subnet (str): Подсеть в формате CIDR (например, "192.168.1.0/24").

        Returns:
            Optional[Device]: Объект локального устройства или None, если не найдено.
        """
        now = datetime.now().astimezone()
        formatted_now = now.strftime("%d-%m-%Y, %H:%M")
        try:
            subnet_network = ipaddress.ip_network(subnet, strict=False)
            interfaces = netifaces.interfaces()
            for iface in interfaces:
                if iface == "lo":
                    continue
                addrs = netifaces.ifaddresses(iface)
                if netifaces.AF_INET in addrs and netifaces.AF_LINK in addrs:
                    ip = addrs[netifaces.AF_INET][0]["addr"]
                    if ipaddress.ip_address(ip) in subnet_network:
                        mac = addrs[netifaces.AF_LINK][0]["addr"]
                        vendor = self.device_info.get_vendor(mac)
                        hostname = self.device_info.get_hostname(ip)
                        open_ports = self.port_scanner.scan(ip)
                        logger.debug(
                            f"Local device: IP={ip}, MAC={mac}, Ports={open_ports}, Hostname={hostname}"
                        )
                        return Device(
                            ip=ip,
                            mac=mac,
                            hostname=hostname,
                            vendor=vendor,
                            open_ports=open_ports,
                            status="online",
                            last_seen=formatted_now,
                            note=None,
                        )
            logger.warning(
                f"Local device not found in subnet {subnet}")
            return None
        except Exception as e:
            logger.warning(
                f"Failed to get local device: {str(e)}")
            return None

    def add_note(self, mac: str, note: str) -> Tuple[bool, Optional[str]]:
        """Добавляет заметку к устройству по MAC-адресу.

        Args:
            mac (str): MAC-адрес устройства.
            note (str): Текст заметки.

        Returns:
            Tuple[bool, Optional[str]]: Кортеж (успех, сообщение об ошибке).
                - Успех: (True, None).
                - Ошибка: (False, сообщение об ошибке).

        Raises:
            Exception: Если произошла ошибка при сохранении.
        """
        try:
            saved_devices = self.storage.load_devices()
            if mac not in saved_devices:
                logger.warning(f"Device with MAC {mac} not found")
                return False, self.translations.gettext("Устройство с MAC {mac} не найдено").format(mac=mac)
            if saved_devices[mac].note is not None:
                logger.error(
                    f"Note for device with MAC {mac} already exists")
                return False, self.translations.gettext("Заметка для устройства с MAC {mac} уже существует").format(mac=mac)
            saved_devices[mac].note = note
            self.storage.save_devices(list(saved_devices.values()))
            logger.info(
                f"Added note '{note}' for device with MAC {mac}")
            return True, None
        except Exception as e:
            logger.error(
                f"Error adding note for MAC {mac}: {str(e)}")
            return False, self.translations.gettext("Внутренняя ошибка сервера: {error}").format(error=str(e))

    def update_note(self, mac: str, note: str) -> Tuple[bool, Optional[str]]:
        """Обновляет заметку устройства по MAC-адресу.

        Args:
            mac (str): MAC-адрес устройства.
            note (str): Новый текст заметки.

        Returns:
            Tuple[bool, Optional[str]]: Кортеж (успех, сообщение об ошибке).
                - Успех: (True, None).
                - Ошибка: (False, сообщение об ошибке).

        Raises:
            Exception: Если произошла ошибка при сохранении.
        """
        try:
            saved_devices = self.storage.load_devices()
            if mac not in saved_devices:
                logger.warning(f"Device with MAC {mac} not found")
                return False, self.translations.gettext("Устройство с MAC {mac} не найдено").format(mac=mac)
            if saved_devices[mac].note is None:
                logger.error(
                    f"Note for device with MAC {mac} does not exist")
                return False, self.translations.gettext("Заметка для устройства с MAC {mac} не существует").format(mac=mac)
            saved_devices[mac].note = note
            self.storage.save_devices(list(saved_devices.values()))
            logger.info(
                f"Updated note '{note}' for device with MAC {mac}")
            return True, None
        except Exception as e:
            logger.error(
                f"Error updating note for MAC {mac}: {str(e)}")
            return False, self.translations.gettext("Внутренняя ошибка сервера: {error}").format(error=str(e))

    def delete_note(self, mac: str) -> Tuple[bool, Optional[str]]:
        """Удаляет заметку устройства по MAC-адресу.

        Args:
            mac (str): MAC-адрес устройства.

        Returns:
            Tuple[bool, Optional[str]]: Кортеж (успех, сообщение об ошибке).
                - Успех: (True, None).
                - Ошибка: (False, сообщение об ошибке).

        Raises:
            Exception: Если произошла ошибка при сохранении.
        """
        try:
            saved_devices = self.storage.load_devices()
            if mac not in saved_devices:
                logger.warning(f"Device with MAC {mac} not found")
                return False, self.translations.gettext("Устройство с MAC {mac} не найдено").format(mac=mac)
            if saved_devices[mac].note is None:
                logger.error(
                    f"Note for device with MAC {mac} does not exist")
                return False, self.translations.gettext("Заметка для устройства с MAC {mac} не существует").format(mac=mac)
            saved_devices[mac].note = None
            self.storage.save_devices(list(saved_devices.values()))
            logger.info(f"Deleted note for device with MAC {mac}")
            return True, None
        except Exception as e:
            logger.error(
                f"Error deleting note for MAC {mac}: {str(e)}")
            return False, self.translations.gettext("Внутренняя ошибка сервера: {error}").format(error=str(e))