gonets/SRV/services/wifi_service.py

import time
import logging
import subprocess
import re
from functools import lru_cache
from typing import List, Optional, Dict, Tuple
from mac_vendor_lookup import MacLookup
import pywifi
from models.wifi import WifiAccessPoint, WifiScanResult
from pathlib import Path
import json

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

class WifiStorage:
    """Класс для хранения данных о Wi-Fi сетях и их заметках по SSID.

    Предоставляет методы для загрузки и сохранения заметок, связанных с SSID сетей,
    в JSON-файле на диске.
    """

    def __init__(self, storage_path: str = "storage/wifi/wifi.json"):
        """Инициализирует объект хранилища Wi-Fi с указанным путём к файлу.

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

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

    def load_notes(self) -> Dict[str, Optional[str]]:
        """Загружает заметки по SSID из файла хранения.

        Returns:
            Dict[str, Optional[str]]: Словарь, где ключ — SSID, значение — заметка
                или None, если заметки нет.

        Raises:
            Exception: Если произошла ошибка при чтении или парсинге файла.
        """
        try:
            if self.storage_file.exists():
                with open(self.storage_file, "r", encoding="utf-8") as f:
                    data = json.load(f)
                    return {item["ssid"]: item.get("note") for item in data.get("networks", [])}
            return {}
        except Exception as e:
            logger.error(f"Ошибка при загрузке заметок: {str(e)}")
            return {}

    def save_notes(self, notes: Dict[str, Optional[str]]) -> None:
        """Сохраняет заметки по SSID в файл.

        Args:
            notes (Dict[str, Optional[str]]): Словарь с SSID как ключами и заметками
                как значениями.

        Raises:
            Exception: Если произошла ошибка при записи в файл.
        """
        try:
            data = {"networks": [{"ssid": ssid, "note": note} for ssid, note in notes.items()]}
            with open(self.storage_file, "w", encoding="utf-8") as f:
                json.dump(data, f, indent=2, ensure_ascii=False)
            logger.info(f"Сохранено {len(notes)} заметок в {self.storage_file}")
        except Exception as e:
            logger.error(f"Ошибка при сохранении заметок: {str(e)}")

class WifiService:
    """Сервис для сканирования текущих Wi-Fi точек доступа.

    Предоставляет функционал для обнаружения Wi-Fi сетей с использованием pywifi и iw,
    определения характеристик точек доступа (сигнал, шифрование, канал и т.д.),
    а также управления заметками.
    """

    def __init__(self):
        """Инициализирует сервис Wi-Fi с настройкой интерфейса и хранилища.

        Инициализирует объект MacLookup для определения вендоров по MAC-адресам,
        PyWiFi для работы с Wi-Fi-интерфейсами и объект хранилища заметок.
        """
        self.mac_lookup = MacLookup()
        self.wifi = pywifi.PyWiFi()
        self.interface = self._get_wifi_interface()
        self.storage = WifiStorage()

    def _get_wifi_interface(self) -> Optional[object]:
        """Автоматически определяет активный Wi-Fi-интерфейс.

        Returns:
            Optional[object]: Объект интерфейса PyWiFi или None, если интерфейс не найден.

        Проверяет доступные интерфейсы, выполняет тестовое сканирование и возвращает
        первый активный интерфейс, исключая P2P-интерфейсы.
        """
        try:
            interfaces = self.wifi.interfaces()
            if not interfaces:
                logger.error("Не найдено сетевых интерфейсов")
                return None

            logger.debug(
                f"Доступные интерфейсы: {[iface.name() for iface in interfaces]}"
            )

            for iface in interfaces:
                iface_name = iface.name()
                if "p2p" in iface_name.lower():
                    logger.debug(f"Пропущен интерфейс {iface_name} (P2P)")
                    continue
                try:
                    iface.scan()
                    time.sleep(1)
                    results = iface.scan_results()
                    if results:
                        logger.info(f"Выбран Wi-Fi-интерфейс: {iface_name}")
                        return iface
                    logger.debug(
                        f"Интерфейс {iface_name} не вернул результатов сканирования"
                    )
                except Exception as e:
                    logger.debug(f"Интерфейс {iface_name} неактивен: {str(e)}")

            logger.error("Не найден активный Wi-Fi-интерфейс")
            return None
        except Exception as e:
            logger.error(f"Ошибка при определении Wi-Fi-интерфейса: {str(e)}")
            return None

    @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"Не удалось определить производителя для MAC {mac}: {str(e)}")
            return "Unknown"

    def _get_encryption_from_iw(self, ap_info: Dict) -> str:
        """Определяет тип шифрования из данных, полученных через iw.

        Args:
            ap_info (Dict): Словарь с данными точки доступа, включая поле "security".

        Returns:
            str: Тип шифрования ("WPA3", "WPA2", "WPA", "WEP" или "Open").
        """
        security = ap_info.get("security", "")
        logger.debug(
            f"Шифрование для BSSID {ap_info.get('bssid')}: {security}")
        if "WPA3" in security:
            return "WPA3"
        elif "WPA2" in security or "RSN" in security:
            return "WPA2"
        elif "WPA" in security:
            return "WPA"
        elif "WEP" in security:
            return "WEP"
        return "Open"

    def _get_channel(self, frequency: Optional[int]) -> Optional[int]:
        """Преобразует частоту в номер канала.

        Args:
            frequency (Optional[int]): Частота в МГц или None.

        Returns:
            Optional[int]: Номер канала или None, если частота вне диапазона.
        """
        if frequency is None:
            return None
        try:
            if 2412 <= frequency <= 2472:
                return (frequency - 2412) // 5 + 1
            elif 5180 <= frequency <= 5825:
                return (frequency - 5180) // 5 + 36
            logger.debug(f"Частота {frequency} вне диапазона 2.4/5 GHz")
            return None
        except Exception as e:
            logger.debug(
                f"Ошибка при определении канала для частоты {frequency}: {str(e)}"
            )
            return None

    def _get_frequency(self, frequency: Optional[int]) -> Optional[str]:
        """Определяет частотный диапазон на основе частоты.

        Args:
            frequency (Optional[int]): Частота в МГц или None.

        Returns:
            Optional[str]: Диапазон ("2.4 GHz", "5 GHz" или None).
        """
        if frequency is None:
            return None
        try:
            if 2412 <= frequency <= 2472:
                return "2.4 GHz"
            elif 5180 <= frequency <= 5825:
                return "5 GHz"
            logger.debug(f"Частота {frequency} вне диапазона 2.4/5 GHz")
            return None
        except Exception as e:
            logger.debug(
                f"Ошибка при определении диапазона для частоты {frequency}: {str(e)}"
            )
            return None

    def _get_iw_scan_data(self) -> List[Dict]:
        """Получает данные сканирования Wi-Fi через команду iw.

        Returns:
            List[Dict]: Список словарей с данными точек доступа (BSSID, SSID, частота,
                сигнал, шифрование).

        Raises:
            subprocess.CalledProcessError: Если команда iw завершилась с ошибкой.
            subprocess.TimeoutExpired: Если превышен таймаут выполнения.
            Exception: Если произошла ошибка при парсинге вывода.
        """
        try:
            if not self.interface:
                logger.error("Интерфейс не определён")
                return []
            cmd = ["sudo", "iw", "dev", self.interface.name(), "scan"]
            logger.debug(f"Выполняем команду: {' '.join(cmd)}")
            result = subprocess.run(cmd,
                                    capture_output=True,
                                    text=True,
                                    timeout=10)
            if result.returncode != 0:
                logger.error(f"Ошибка iw scan: {result.stderr}")
                return []
            output = result.stdout
            logger.debug(
                f"Вывод iw scan (первые 500 символов): {output[:500]}...")

            aps = []
            current_ap = None
            bssid_pattern = re.compile(r"BSS ([0-9a-f:]{17})")
            ssid_pattern = re.compile(r"SSID: ?(.*)")
            freq_pattern = re.compile(r"freq: (\d+)")
            signal_pattern = re.compile(r"signal: ([-\d.]+) dBm")
            rsn_pattern = re.compile(r"RSN:")
            wpa_pattern = re.compile(r"WPA:")
            security_type_pattern = re.compile(r"\* (WPA3|WPA2|WPA|WEP)")

            for line in output.splitlines():
                line = line.strip()
                if bssid_match := bssid_pattern.match(line):
                    if current_ap:
                        aps.append(current_ap)
                    current_ap = {
                        "bssid": bssid_match.group(1).lower(),
                        "security": ""
                    }
                elif current_ap:
                    if ssid_match := ssid_pattern.match(line):
                        current_ap["ssid"] = (ssid_match.group(1)
                                              if ssid_match.group(1).strip()
                                              else "")
                    elif freq_match := freq_pattern.match(line):
                        current_ap["frequency"] = int(freq_match.group(1))
                    elif signal_match := signal_pattern.match(line):
                        current_ap["signal"] = float(signal_match.group(1))
                    elif rsn_pattern.match(line):
                        current_ap["security"] += "RSN;"
                    elif wpa_pattern.match(line):
                        current_ap["security"] += "WPA;"
                    elif security_type_match := security_type_pattern.search(
                            line):
                        current_ap["security"] += security_type_match.group(
                            1) + ";"

            if current_ap:
                aps.append(current_ap)

            logger.debug(
                f"Найдено {len(aps)} точек через iw: {[ap['bssid'] for ap in aps]}"
            )
            for ap in aps:
                logger.debug(f"Данные точки {ap['bssid']}: {ap}")
            return aps
        except subprocess.CalledProcessError as e:
            logger.error(f"Ошибка при выполнении iw scan: {e.stderr}")
            return []
        except subprocess.TimeoutExpired:
            logger.error("Таймаут при выполнении iw scan")
            return []
        except Exception as e:
            logger.error(f"Ошибка при парсинге iw scan: {str(e)}")
            return []

    def scan_wifi(self, timeout: int = 30) -> WifiScanResult:
        """Сканирует текущие Wi-Fi точки доступа.

        Args:
            timeout (int, optional): Таймаут ожидания в секундах. По умолчанию 30.

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

        Сканирует Wi-Fi с использованием pywifi и iw, объединяет данные, определяет
        лучшие точки по сигналу и сохраняет заметки.
        """
        start_time = time.time()
        aps = []

        if not self.interface:
            logger.error("Wi-Fi-интерфейс не доступен")
            return WifiScanResult(access_points=aps)

        try:
            logger.info(
                f"Начало сканирования Wi-Fi на интерфейсе {self.interface.name()}"
            )
            self.interface.scan()
            time.sleep(10)
            py_results = self.interface.scan_results()
            logger.debug(f"Найдено {len(py_results)} точек через pywifi")

            iw_results = self._get_iw_scan_data()
            iw_by_bssid = {ap["bssid"]: ap for ap in iw_results if ap.get("bssid")}

            best_aps_by_ssid: Dict[str, WifiAccessPoint] = {}
            saved_notes = self.storage.load_notes()

            for ap in py_results:
                bssid = ap.bssid.lower() if ap.bssid else None
                if not bssid:
                    logger.debug("Пропущена точка без BSSID")
                    continue
                ssid = ap.ssid if ap.ssid else ""
                hidden = not ssid
                signal = ap.signal if ap.signal else None
                iw_data = iw_by_bssid.get(bssid, {})

                frequency = iw_data.get("frequency")
                channel = self._get_channel(frequency)
                freq_band = self._get_frequency(frequency)
                encryption = self._get_encryption_from_iw(iw_data)
                vendor = self.get_vendor(bssid)

                logger.debug(
                    f"Точка {bssid}: freq={frequency}, channel={channel}, encryption={encryption}, signal={signal}"
                )

                current_ap = WifiAccessPoint(
                    ssid=ssid,
                    bssid=bssid,
                    channel=channel,
                    signal=signal if signal else iw_data.get("signal"),
                    encryption=encryption,
                    frequency=freq_band,
                    vendor=vendor,
                    hidden=hidden,
                )

                if ssid not in best_aps_by_ssid or (
                        current_ap.signal and best_aps_by_ssid[ssid].signal
                        and current_ap.signal > best_aps_by_ssid[ssid].signal):
                    current_ap.note = saved_notes.get(ssid)
                    best_aps_by_ssid[ssid] = current_ap

            aps = list(best_aps_by_ssid.values())
            notes_to_save = {ap.ssid: ap.note for ap in aps if ap.ssid}
            self.storage.save_notes(notes_to_save)

            logger.info(
                f"Сканирование Wi-Fi завершено за {time.time() - start_time:.2f} секунд, "
                f"найдено {len(aps)} уникальных сетей")
            return WifiScanResult(access_points=aps)

        except Exception as e:
            logger.error(f"Ошибка при сканировании Wi-Fi: {str(e)}")
            return WifiScanResult(access_points=aps)

    def add_note(self, ssid: str, note: str) -> Tuple[bool, Optional[str]]:
        """Добавляет заметку к сети Wi-Fi по SSID.

        Args:
            ssid (str): SSID сети Wi-Fi.
            note (str): Текст заметки.

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

        Raises:
            Exception: Если произошла ошибка при сохранении заметки.
        """
        try:
            saved_notes = self.storage.load_notes()
            if ssid not in saved_notes or saved_notes[ssid] is None:
                logger.info(f"Добавлена новая заметка для сети с SSID {ssid}")
            elif saved_notes[ssid] is not None:
                logger.error(f"Заметка для сети с SSID {ssid} уже существует")
                return False, f"Заметка для сети с SSID {ssid} уже существует"
            saved_notes[ssid] = note
            self.storage.save_notes(saved_notes)
            logger.info(f"Добавлена заметка '{note}' для сети с SSID {ssid}")
            return True, None
        except Exception as e:
            logger.error(f"Ошибка при добавлении заметки для SSID {ssid}: {str(e)}")
            return False, f"Внутренняя ошибка сервера: {str(e)}"

    def update_note(self, ssid: str, note: str) -> Tuple[bool, Optional[str]]:
        """Обновляет заметку сети Wi-Fi по SSID.

        Args:
            ssid (str): SSID сети Wi-Fi.
            note (str): Новый текст заметки.

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

        Raises:
            Exception: Если произошла ошибка при сохранении заметки.
        """
        try:
            saved_notes = self.storage.load_notes()
            if ssid not in saved_notes:
                logger.warning(f"Сеть с SSID {ssid} не найдена")
                return False, f"Сеть с SSID {ssid} не найдена"
            if saved_notes[ssid] is None:
                logger.error(f"Заметка для сети с SSID {ssid} не существует")
                return False, f"Заметка для сети с SSID {ssid} не существует"
            saved_notes[ssid] = note
            self.storage.save_notes(saved_notes)
            logger.info(f"Обновлена заметка '{note}' для сети с SSID {ssid}")
            return True, None
        except Exception as e:
            logger.error(f"Ошибка при обновлении заметки для SSID {ssid}: {str(e)}")
            return False, f"Внутренняя ошибка сервера: {str(e)}"

    def delete_note(self, ssid: str) -> Tuple[bool, Optional[str]]:
        """Удаляет заметку сети Wi-Fi по SSID.

        Args:
            ssid (str): SSID сети Wi-Fi.

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

        Raises:
            Exception: Если произошла ошибка при сохранении изменений.
        """
        try:
            saved_notes = self.storage.load_notes()
            if ssid not in saved_notes:
                logger.warning(f"Сеть с SSID {ssid} не найдена")
                return False, f"Сеть с SSID {ssid} не найдена"
            if saved_notes[ssid] is None:
                logger.error(f"Заметка для сети с SSID {ssid} не существует")
                return False, f"Заметка для сети с SSID {ssid} не существует"
            saved_notes[ssid] = None
            self.storage.save_notes(saved_notes)
            logger.info(f"Удалена заметка для сети с SSID {ssid}")
            return True, None
        except Exception as e:
            logger.error(f"Ошибка при удалении заметки для SSID {ssid}: {str(e)}")
            return False, f"Внутренняя ошибка сервера: {str(e)}"