ok
Direktori : /opt/imunify360/venv/lib64/python3.11/site-packages/im360/model/ |
Current File : //opt/imunify360/venv/lib64/python3.11/site-packages/im360/model/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 )