22

✏️ Editing: port_ips_deny_mode.py

import ipaddress
import logging
from pathlib import Path
from typing import List, Literal, Union

logger = logging.getLogger(__name__)

PROTOS = (TCP, UDP, ALL) = ("tcp", "udp", "all")

class IpPortDenyModeError(Exception):
    pass

Proto = Literal["tcp", "udp", "all"]
PortNetworkMap = dict[Proto, dict[int, List[str]]]

class WhitelistPortIPsDenyMode:
    _LIST_PATH = "/etc/imunify360/whitelist/ports"

@classmethod
    def load(cls) -> PortNetworkMap:
        """Read port:networks from the file"""
        port_to_networks = {TCP: {}, UDP: {}}
        for path in Path(cls._LIST_PATH).glob("*.txt"):
            port_to_networks = cls._load_file(
                path.absolute(), port_to_networks
            )
        return port_to_networks

@classmethod
    def _load_file(
        cls, path: Path, port_to_networks: PortNetworkMap
    ) -> PortNetworkMap:
        for line_num, line in enumerate(path.open().readlines()):
            item = line.partition("#")[0].strip()
            if not item:
                continue

network_list = []
            valid = True

port, semicolon, tail = item.partition(":")

protos: List[Proto] = [TCP, UDP]
            proto, _, networks = tail.partition(":")
            proto: Proto = proto.strip().lower()
            if proto in PROTOS:
                if proto != ALL:
                    protos = [proto]
            else:
                networks = tail

if not semicolon:
                valid = False

try:
                port = int(port)
            except ValueError:
                valid = False
            else:
                for network in networks.split(","):
                    network = network.strip()
                    if network:
                        try:
                            ipaddress.ip_network(network)
                            network_list.append(network)
                        except ValueError:
                            valid = False

for proto in protos:
                    if network_list:
                        existing_ports = port_to_networks[proto].setdefault(
                            port, []
                        )
                        existing_ports.extend(network_list)
                        port_to_networks[proto][port] = sorted(
                            list(set(existing_ports))
                        )

if not valid:
                logger.error(
                    "Wrong IP/Subnet format in %s:%s. Use CIDR format.",
                    path,
                    line_num,
                )

return port_to_networks

@classmethod
    def count(
        cls,
        proto: Union[
            Literal["tcp", "udp", "all"], List[Literal["tcp", "udp"]]
        ],
    ) -> int:
        match proto:
            case _ if isinstance(proto, list) and all(
                p in ["tcp", "udp"] for p in proto
            ):
                protos = proto
            case "tcp" | "udp":
                protos = [proto]
            case "all":
                protos = ["tcp", "udp"]
            case _:
                raise ValueError("Invalid input")
        proto_to_ports = cls.load()
        return sum(
            sum(len(networks) for networks in proto_to_ports[p].values())
            for p in protos
        )

MR SKB