import ipaddress import subprocess from typing import List, Union from models.ufw import Direction, Protocol, UFWRule, UFWRulesResponse, UFWStatus DEBUG = True class UFWService: """ Сервис для взаимодействия с утилитой UFW (Uncomplicated Firewall) в Linux. Позволяет получить список правил UFW и управлять ими. """ def _run_command(self, command: str) -> str: """ Выполняет команду в системе и возвращает ее вывод. Args: command (str): Команда для выполнения в системе. Returns: str: Вывод команды. """ try: if DEBUG: print(f"Running command: {command}") # Используем shell=True для корректной обработки аргументов с пробелами result = subprocess.run( f"sudo {command}", shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, ) if DEBUG: print(f"Command output: {result.stdout}") print(f"Command error: {result.stderr}") if result.returncode != 0: raise subprocess.CalledProcessError( result.returncode, command, result.stdout, result.stderr ) return result.stdout except subprocess.CalledProcessError as e: if DEBUG: print(f"Command failed with return code {e.returncode}") print(f"Error output: {e.stderr}") raise Exception(f"Command failed: {e.stderr}") def _get_port_range_args(self, port: Union[int, str]) -> List[str]: """ Преобразует порт или диапазон портов в аргументы команды UFW. Args: port: Порт (число) или диапазон портов (строка вида "3000:3999") Returns: List[str]: Список аргументов для UFW """ if isinstance(port, int): return [str(port)] elif isinstance(port, str) and ":" in port: # Диапазон портов уже в нужном формате return [port] else: raise ValueError(f"Некорректный формат порта: {port}") def get_ufw_rules(self) -> UFWRulesResponse: """ Получает список правил UFW. Returns: UFWRulesResponse: Объект, содержащий статус UFW и список правил """ # Check UFW status status_result = self._run_command("ufw status numbered") lines = status_result.splitlines() if not lines: return UFWRulesResponse(enabled=False, rules=[]) first_line = lines[0] is_enabled = "Status: active" in first_line if not is_enabled: return UFWRulesResponse(enabled=False, rules=[]) rules = [] # Skip header lines for line in lines[4:]: # Skip Status and header lines if not line.strip() or not line.startswith("["): continue try: rule = UFWRule.from_ufw_output(line) rules.append(rule) except Exception as e: if DEBUG: print(f"Ошибка парсинга правила: {line}, ошибка: {str(e)}") continue return UFWRulesResponse(enabled=is_enabled, rules=rules) def get_ufw_status(self) -> UFWStatus: """ Проверяет статус UFW. Returns: UFWStatus: Объект с информацией о статусе UFW """ try: status_result = self._run_command("ufw status") first_line = ( status_result.splitlines()[0] if status_result.splitlines() else "" ) is_enabled = "Status: active" in first_line message = "UFW активен" if is_enabled else "UFW неактивен" if DEBUG: print(f"Результат статуса: {status_result}") print(f"Включен: {is_enabled}") return UFWStatus(enabled=is_enabled, message=message) except Exception as e: if DEBUG: print(f"Ошибка при проверке статуса UFW: {e}") return UFWStatus( enabled=False, message=f"Ошибка при проверке статуса: {str(e)}", ) def enable_ufw(self) -> bool: """ Включает UFW. Returns: bool: True если успешно включен, False в противном случае """ try: self._run_command("ufw enable") return True except Exception as e: if DEBUG: print(f"Ошибка при включении UFW: {e}") return False def disable_ufw(self) -> bool: """ Выключает UFW. Returns: bool: True если успешно выключен, False в противном случае """ try: self._run_command("ufw disable") return True except Exception as e: if DEBUG: print(f"Ошибка при выключении UFW: {e}") return False def add_rules_base(self, rules: List[UFWRule], direction: str = "") -> bool: """ Базовый метод для добавления правил UFW. Args: rules: Список правил для добавления direction: Направление трафика ("in" или "out") Returns: bool: True если все правила были успешно добавлены, False в противном случае """ try: for rule in rules: # Устанавливаем направление правила rule.direction = Direction.IN if direction == "in" else Direction.OUT # Формируем базовую команду command = f"ufw {rule.action.value}" # Добавляем направление command += f" {rule.direction.value}" # Добавляем from IP если он указан if rule.from_ip and rule.from_ip != "any": command += f" from {rule.from_ip}" else: command += " from any" # Добавляем to IP если он указан if rule.to_ip and rule.to_ip != "any": command += f" to {rule.to_ip}" else: command += " to any" # Добавляем порт если он указан if rule.port: command += f" port {rule.port}" # Добавляем протокол if rule.protocol != Protocol.BOTH: command += f" proto {rule.protocol.value}" # Добавляем комментарий если он есть if rule.comment: command += f" comment '{rule.comment}'" self._run_command(command) return True except Exception as e: if DEBUG: print(f"Error adding UFW rules: {str(e)}") return False def add_ufw_in_rules(self, rules: List[UFWRule]) -> bool: """ Добавляет входящие (in) правила UFW. Args: rules (List[UFWRule]): Список правил для добавления Returns: bool: True если все правила были успешно добавлены, False в противном случае """ return self.add_rules_base(rules, "in") def add_ufw_out_rules(self, rules: List[UFWRule]) -> bool: """ Добавляет исходящие (out) правила UFW. Args: rules (List[UFWRule]): Список правил для добавления Returns: bool: True если все правила были успешно добавлены, False в противном случае """ return self.add_rules_base(rules, "out") def add_rules(self, rules: List[UFWRule]) -> dict: """ Добавляет правила UFW. Args: rules: Список правил для добавления Returns: dict: Результат добавления правил """ results = [] for rule in rules: try: # Определяем версию IP from_ip = ipaddress.ip_network(rule.from_ip, strict=False) is_ipv6 = isinstance(from_ip, ipaddress.IPv6Network) # Получаем список протоколов protocols = rule.protocol.protocols print(f"Обработка правила для протоколов: {protocols}") for proto in protocols: # Формируем базовую команду command = f"ufw {rule.action.value}" # Добавляем IP адреса command += f" from {rule.from_ip}" if rule.to_ip: command += f" to {rule.to_ip}" # Добавляем порт и протокол port_args = self._get_port_range_args(rule.port) command += f" port {' '.join(port_args)} proto {proto}" # Добавляем флаг IPv6 если нужно if is_ipv6: command += " v6" # Добавляем комментарий если есть if rule.comment: command += f" comment '{rule.comment}'" print(f"Подготовлена команда: {command}") # Выполняем команду try: output = self._run_command(command) results.append({"success": True, "message": output}) except Exception as e: error_msg = f"Ошибка выполнения команды: {str(e)}" print(error_msg) results.append({"success": False, "message": error_msg}) except Exception as e: error_msg = f"Ошибка при добавлении правила: {str(e)}" print(error_msg) results.append({"success": False, "message": error_msg}) # Выводим результаты print("\nРезультаты добавления правил:") for result in results: status = "Успешно" if result["success"] else "Ошибка" print(f" - {status}: {result['message']}") return {"success": True, "results": results} def delete_rules(self, rule_numbers: List[int]) -> dict: """ Удаляет правила UFW по их номерам. Args: rule_numbers (List[int]): Список номеров правил для удаления Returns: dict: Результат операции удаления """ try: # Проверяем существование правил status_result = self._run_command("ufw status numbered") lines = status_result.splitlines() # Пропускаем заголовок (первые 4 строки) rules = [] for line in lines[4:]: # Skip header lines if not line.strip() or not line.startswith("["): continue try: # Извлекаем номер правила из строки [N] rule_num = int(line[1 : line.index("]")]) rules.append(rule_num) except (ValueError, IndexError): continue if DEBUG: print(f"Найдены правила: {rules}") results = [] # Сортируем номера в обратном порядке, чтобы удаление происходило с конца # Это предотвратит смещение номеров при удалении for rule_number in sorted(rule_numbers, reverse=True): try: # Проверяем существует ли правило с таким номером if rule_number not in rules: results.append( { "number": rule_number, "success": False, "error": f"Правило с номером {rule_number} не найдено", } ) continue # Удаляем правило self._run_command(f"ufw --force delete {rule_number}") results.append( { "number": rule_number, "success": True, "message": f"Правило {rule_number} успешно удалено", } ) except Exception as e: if DEBUG: print(f"Ошибка при удалении правила {rule_number}: {e}") results.append( {"number": rule_number, "success": False, "error": str(e)} ) # Подсчитываем статистику total = len(rule_numbers) successful = len([r for r in results if r["success"]]) return { "success": True, "total": total, "successful": successful, "results": results, } except Exception as e: if DEBUG: print(f"Ошибка при удалении правил: {e}") return {"success": False, "error": str(e), "results": []}