409 lines
16 KiB
Python
409 lines
16 KiB
Python
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).
|
||
|
||
Назначение:
|
||
- Чтение статуса и списка правил брандмауэра.
|
||
- Добавление, удаление и массовая обработка правил с поддержкой диапазонов портов,
|
||
комментариев и выбором протоколов.
|
||
|
||
Детали реализации:
|
||
- Команды выполняются через `sudo ufw ...`; вывод парсится в модели pydantic.
|
||
- Для операций удаления соблюдается порядок с конца, чтобы не нарушать нумерацию.
|
||
"""
|
||
|
||
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": []}
|