283 lines
9.6 KiB
Python
283 lines
9.6 KiB
Python
from typing import List, Optional, Annotated, Union
|
||
from enum import Enum
|
||
from pydantic import BaseModel, Field, validator
|
||
import ipaddress
|
||
|
||
|
||
class Protocol(str, Enum):
|
||
"""
|
||
Enum для сетевых протоколов UFW.
|
||
"""
|
||
|
||
TCP = "tcp"
|
||
UDP = "udp"
|
||
BOTH = "both"
|
||
|
||
@property
|
||
def protocols(self) -> List[str]:
|
||
"""Возвращает список протоколов для использования в команде UFW."""
|
||
return ["tcp", "udp"] if self == Protocol.BOTH else [self.value]
|
||
|
||
|
||
class Action(str, Enum):
|
||
"""
|
||
Enum для действий UFW.
|
||
|
||
Attributes:
|
||
ALLOW: Разрешить трафик
|
||
DENY: Запретить трафик без уведомления
|
||
REJECT: Запретить трафик с уведомлением отправителю
|
||
"""
|
||
|
||
ALLOW = "allow"
|
||
DENY = "deny"
|
||
REJECT = "reject"
|
||
|
||
|
||
class Direction(str, Enum):
|
||
"""
|
||
Enum для направления трафика UFW.
|
||
|
||
Attributes:
|
||
IN: Входящий трафик
|
||
OUT: Исходящий трафик
|
||
"""
|
||
|
||
IN = "in"
|
||
OUT = "out"
|
||
|
||
|
||
class UFWRule(BaseModel):
|
||
"""Модель правила UFW"""
|
||
|
||
action: Action
|
||
direction: Direction
|
||
from_ip: str
|
||
to_ip: str
|
||
port: Union[int, str]
|
||
protocol: Protocol
|
||
comment: Optional[str] = None
|
||
|
||
@validator("port")
|
||
def validate_port(cls, v):
|
||
"""
|
||
Проверяет корректность порта или диапазона портов.
|
||
|
||
Args:
|
||
v: Порт (число) или диапазон портов (строка вида "3000:3999")
|
||
|
||
Returns:
|
||
Union[int, str]: Валидированное значение порта
|
||
|
||
Raises:
|
||
ValueError: Если порт или диапазон портов некорректен
|
||
"""
|
||
if isinstance(v, int):
|
||
if not (1 <= v <= 65535):
|
||
raise ValueError(f"Порт должен быть в диапазоне 1-65535, получено: {v}")
|
||
return v
|
||
elif isinstance(v, str):
|
||
try:
|
||
start, end = map(int, v.split(":"))
|
||
if not (1 <= start <= end <= 65535):
|
||
raise ValueError
|
||
return v
|
||
except (ValueError, TypeError):
|
||
raise ValueError(
|
||
f"Некорректный диапазон портов: {v}. Формат: start:end, где 1 <= start <= end <= 65535"
|
||
)
|
||
else:
|
||
raise ValueError(
|
||
f"Порт должен быть числом или строкой с диапазоном, получено: {type(v)}"
|
||
)
|
||
|
||
@validator("from_ip", "to_ip")
|
||
def validate_ip(cls, v):
|
||
"""Проверяет корректность IP адреса"""
|
||
try:
|
||
# Пробуем создать объект IP-сети, что автоматически проверит корректность адреса
|
||
ipaddress.ip_network(v, strict=False)
|
||
return v
|
||
except ValueError:
|
||
raise ValueError(f"Некорректный IP адрес или подсеть: {v}")
|
||
|
||
@validator("to_ip")
|
||
def validate_ip_version_compatibility(cls, v, values):
|
||
"""Проверяет совместимость версий IP адресов"""
|
||
if "from_ip" not in values:
|
||
return v
|
||
|
||
# Определяем версии IP
|
||
from_ip = ipaddress.ip_network(values["from_ip"], strict=False)
|
||
to_ip = ipaddress.ip_network(v, strict=False)
|
||
|
||
if isinstance(from_ip, ipaddress.IPv6Network) != isinstance(
|
||
to_ip, ipaddress.IPv6Network
|
||
):
|
||
raise ValueError(
|
||
"Несовместимые версии IP: нельзя смешивать IPv4 и IPv6 в одном правиле"
|
||
)
|
||
|
||
return v
|
||
|
||
@validator("action", pre=True)
|
||
def validate_action(cls, v):
|
||
"""Приводит значение action к нижнему регистру перед валидацией"""
|
||
if isinstance(v, str):
|
||
return v.lower()
|
||
return v
|
||
|
||
@validator("protocol", pre=True)
|
||
def validate_protocol(cls, v):
|
||
"""Приводит значение protocol к нижнему регистру перед валидацией"""
|
||
if isinstance(v, str):
|
||
return v.lower()
|
||
return v
|
||
|
||
def to_ufw_command(self) -> str:
|
||
"""Преобразует правило в команду UFW."""
|
||
cmd_parts = ["ufw"]
|
||
|
||
# Добавляем действие
|
||
cmd_parts.append(self.action.value)
|
||
|
||
# Добавляем направление
|
||
cmd_parts.append(self.direction.value)
|
||
|
||
# Добавляем from с IP
|
||
cmd_parts.extend(["from", self.from_ip])
|
||
|
||
# Добавляем to с IP
|
||
cmd_parts.extend(["to", self.to_ip])
|
||
|
||
# Добавляем порт и протокол
|
||
port_str = str(self.port) # Работает как с числами, так и с диапазонами
|
||
cmd_parts.extend(["port", port_str])
|
||
|
||
# Добавляем протокол
|
||
proto = self.protocol.value
|
||
if proto == "both":
|
||
cmd_parts.append("proto tcp/udp")
|
||
else:
|
||
cmd_parts.extend(["proto", proto])
|
||
|
||
# Добавляем комментарий если есть
|
||
if self.comment:
|
||
cmd_parts.extend(["comment", f"'{self.comment}'"])
|
||
|
||
return " ".join(cmd_parts)
|
||
|
||
@classmethod
|
||
def from_ufw_output(cls, line: str) -> "UFWRule":
|
||
"""
|
||
Создает экземпляр UFWRule из строки вывода команды 'ufw status numbered'.
|
||
|
||
Args:
|
||
line: Строка вывода UFW (например, "[1] 80/tcp ALLOW IN 192.168.1.0/24 # Web Server")
|
||
|
||
Returns:
|
||
UFWRule: Новый экземпляр правила
|
||
"""
|
||
import re
|
||
|
||
# Регулярное выражение для разбора строки правила UFW
|
||
pattern = r"""^\s*
|
||
\[\s*(?P<number>\d+)\s*\]\s*
|
||
(?P<port>\d+(?::\d+)?)/(?P<protocol>tcp(?:/udp)?|udp)\s+
|
||
(?P<action>ALLOW|DENY|REJECT)\s+(?P<direction>IN|OUT)\s+
|
||
(?P<from_ip>[\w\d\./]+|\bAnywhere\b)
|
||
(?:\s+\(v6\))?
|
||
(?:\s+\((?:out|in)\))?
|
||
\s*
|
||
(?:\#\s*(?P<comment>.*?))?
|
||
\s*$
|
||
"""
|
||
|
||
match = re.match(pattern, line, re.VERBOSE | re.IGNORECASE)
|
||
if not match:
|
||
raise ValueError(f"Невозможно разобрать строку правила UFW: {line}")
|
||
|
||
groups = match.groupdict()
|
||
|
||
# Преобразуем протокол в enum
|
||
proto_str = groups["protocol"].lower()
|
||
protocol = (
|
||
Protocol.BOTH
|
||
if proto_str == "tcp/udp"
|
||
else Protocol.UDP if proto_str == "udp" else Protocol.TCP
|
||
)
|
||
|
||
# Преобразуем действие в enum (игнорируем регистр)
|
||
action_str = groups["action"].lower()
|
||
action = (
|
||
Action.ALLOW
|
||
if action_str == "allow"
|
||
else Action.DENY if action_str == "deny" else Action.REJECT
|
||
)
|
||
|
||
# Преобразуем направление в enum
|
||
direction_str = groups["direction"].lower()
|
||
direction = Direction.IN if direction_str == "in" else Direction.OUT
|
||
|
||
# Обрабатываем IP адреса
|
||
from_ip = groups["from_ip"].strip()
|
||
if from_ip.lower() == "anywhere":
|
||
from_ip = "0.0.0.0/0"
|
||
|
||
# Для исходящих правил меняем местами from и to
|
||
to_ip = "0.0.0.0/0" # По умолчанию anywhere
|
||
|
||
# Обрабатываем порт
|
||
port = groups["port"]
|
||
if ":" in port:
|
||
# Оставляем как строку для диапазона
|
||
port = port.strip()
|
||
else:
|
||
# Преобразуем в число для одиночного порта
|
||
port = int(port)
|
||
|
||
# Обрабатываем комментарий
|
||
comment = groups.get("comment", "").strip() if groups.get("comment") else None
|
||
|
||
return cls(
|
||
action=action,
|
||
direction=direction,
|
||
from_ip=from_ip,
|
||
to_ip=to_ip,
|
||
port=port,
|
||
protocol=protocol,
|
||
comment=comment,
|
||
)
|
||
|
||
|
||
class UFWRulesResponse(BaseModel):
|
||
"""
|
||
Модель для представления списка правил UFW.
|
||
|
||
Attributes:
|
||
enabled: Включен ли UFW
|
||
rules: Список правил UFW
|
||
total_rules: Общее количество правил
|
||
"""
|
||
|
||
enabled: bool = Field(..., description="Включен ли UFW")
|
||
rules: List[UFWRule] = Field(default_factory=list, description="Список правил UFW")
|
||
|
||
@property
|
||
def total_rules(self) -> int:
|
||
"""Возвращает общее количество правил."""
|
||
return len(self.rules)
|
||
|
||
|
||
class UFWStatus(BaseModel):
|
||
"""
|
||
Модель для представления статуса UFW.
|
||
|
||
Attributes:
|
||
enabled: Включен ли UFW
|
||
message: Сообщение о статусе
|
||
last_checked: Время последней проверки статуса
|
||
"""
|
||
|
||
enabled: bool = Field(..., description="Включен ли UFW")
|
||
message: str = Field(..., description="Сообщение о статусе")
|