^^&&^^

Edit File: port.py

import logging
from typing import Dict, FrozenSet, Iterable, List

from defence360agent.contracts.config import PORT_BLOCKING_MODE_ALLOW
from defence360agent.utils import retry_on, timeit
from im360.contracts.config import Firewall
from im360.internals.core import ip_versions
from im360.internals.core.firewall import FirewallRules, get_firewall
from im360.internals.core.firewall.base import FirewallBatchCommandError
from im360.internals.core.ipset import (
    IP_SET_PREFIX,
    AbstractIPSet,
    IPSetAtomicRestoreBase,
    IPSetCount,
    get_ipset_family,
    libipset,
)
from im360.internals.core.ipset.libipset import IPSetCmdBuilder
from im360.model.firewall import BlockedPort, IgnoredByPort
from defence360agent.utils.validate import IP, IPVersion, NumericIPVersion

from .redirect import IPSetNoRedirectPort

logger = logging.getLogger(__name__)

class SingleIPSetPort(IPSetAtomicRestoreBase):
    TEMPLATE = "{prefix}.{ip_version}.ignored-by-{proto}-{port}"

def __init__(self, port, proto):
        super().__init__(port, proto)
        self.port = port
        self.proto = proto

def gen_ipset_name_for_ip_version(self, ip_version: IPVersion) -> str:
        return self.custom_ipset_name or self.TEMPLATE.format(
            prefix=IP_SET_PREFIX,
            ip_version=ip_version,
            proto=self.proto,
            port=self.port,
        )

def gen_ipset_create_ops(self, ip_version: IPVersion) -> List[str]:
        return [
            IPSetCmdBuilder.get_create_cmd(
                self.gen_ipset_name_for_ip_version(ip_version),
                family=get_ipset_family(ip_version),
            )
        ]

def gen_ipset_destroy_ops(self, ip_version: IPVersion) -> List[str]:
        return [
            IPSetCmdBuilder.get_destroy_cmd(
                self.gen_ipset_name_for_ip_version(ip_version)
            )
        ]

def gen_ipset_flush_ops(self, ip_version: IPVersion) -> List[str]:
        return [
            IPSetCmdBuilder.get_flush_cmd(
                self.gen_ipset_name_for_ip_version(ip_version)
            )
        ]

def _fetch(self, ip_version: IPVersion = None):
        result = []
        for row in IgnoredByPort.fetch(
            NumericIPVersion.from_ip_version(ip_version)
        ).where(
            BlockedPort.port == self.port, BlockedPort.proto == self.proto
        ):
            result.append(row.ip)
        return result

def db_count(self, ip_version: IPVersion = None):
        return len(self._fetch(ip_version))

async def gen_ipset_restore_ops(self, ip_version: IPVersion) -> List[str]:
        lines = []
        name = self.gen_ipset_name_for_ip_version(ip_version)
        for ip in self._fetch(ip_version):
            lines.append(f"add {name} {ip} -exist")

return lines

class IPSetPort(AbstractIPSet):
    TCP = "tcp"
    UDP = "udp"
    ALL = "all"

PROTOS = (TCP, UDP, ALL)
    MIN_PORT, MAX_PORT = 0, 65535

def expand_proto(self, proto):
        assert proto in self.PROTOS, 'protocol "{}" is not supported'.format(
            proto
        )

if proto == self.ALL:
            return [self.TCP, self.UDP]
        return [proto]

async def block(self, item, *args, **kwargs):
        port, generic_proto = item

await libipset.create_hash_set(
                    set_name,
                    timeout=0,
                    set_type=libipset.HASH_NET,
                    family=get_ipset_family(ip_version),
                )

actions = [
                    fw.append_rule(
                        FirewallRules.port_rule(
                            set_name,
                            port,
                            proto,
                            policy=FirewallRules.LOG_BLOCK_PORT_CHAIN,
                        ),
                        chain=FirewallRules.BP_INPUT_CHAIN,
                        priority=FirewallRules.DEFAULT_PRIORITY,
                    )
                    for proto in self.expand_proto(generic_proto)
                ]

decorated = retry_on(FirewallBatchCommandError, max_tries=3)(
                    fw.commit
                )
                await decorated(actions)
                await IPSetNoRedirectPort().add_item(port, ip_version)

async def unblock(self, item, *args, **kwargs):
        port, generic_proto = item

for ip_version in ip_versions.enabled():
            async with await get_firewall(ip_version) as fw:
                ipset = SingleIPSetPort(port=port, proto=generic_proto)
                set_name = ipset.gen_ipset_name_for_ip_version(ip_version)
                actions = [
                    fw.delete_rule(
                        FirewallRules.port_rule(
                            set_name,
                            port,
                            proto,
                            policy=FirewallRules.LOG_BLOCK_PORT_CHAIN,
                        ),
                        chain=FirewallRules.BP_INPUT_CHAIN,
                        priority=FirewallRules.DEFAULT_PRIORITY,
                    )
                    for proto in self.expand_proto(generic_proto)
                ]
                decorated = retry_on(FirewallBatchCommandError, max_tries=3)(
                    fw.commit
                )
                await decorated(actions)
                await IPSetNoRedirectPort().delete_item(port, ip_version)

ipset = SingleIPSetPort(port=port, proto=generic_proto)
            set_name = ipset.gen_ipset_name_for_ip_version(ip_version)
            await libipset.delete_set(set_name)

async def gen_ipset_restore_ops(self, ip_version: IPVersion):
        return []

def get_all_ipset_instances(
        self, ip_version: IPVersion
    ) -> List[SingleIPSetPort]:
        if Firewall.port_blocking_mode != PORT_BLOCKING_MODE_ALLOW:
            return []
        result = []
        for port, proto in self._fetch():
            result.append(SingleIPSetPort(port=port, proto=proto))
        return result

def gen_ipset_create_ops(self, ip_version: IPVersion) -> List[str]:
        """
        Generate list of commands to create all ip sets

:return: list of ipset commands to use with ipset restore
        """
        result: List[str] = []
        for ip_set in self.get_all_ipset_instances(ip_version):
            result.extend(ip_set.gen_ipset_create_ops(ip_version))
        return result

def _fetch(self):
        return [(row.port, row.proto) for row in BlockedPort.select()]

def get_all_ipsets(self, ip_version: IPVersion) -> FrozenSet[str]:
        return frozenset(
            ipset.gen_ipset_name_for_ip_version(ip_version)
            for ipset in self.get_all_ipset_instances(ip_version)
        )

def get_rules(self, ip_version: IPVersion, **kwargs) -> Iterable[dict]:
        if Firewall.port_blocking_mode != PORT_BLOCKING_MODE_ALLOW:
            return []
        result = [
            dict(
                rule=FirewallRules.compose_action(
                    FirewallRules.BP_INPUT_CHAIN
                ),
                chain=FirewallRules.IMUNIFY_INPUT_CHAIN,
                table=FirewallRules.FILTER,
                priority=FirewallRules.PORT_PROTO_PRIORITY,
            )
        ]
        for ipset in self.get_all_ipset_instances(ip_version):
            result.extend(
                [
                    dict(
                        rule=FirewallRules.port_rule(
                            ipset.gen_ipset_name_for_ip_version(ip_version),
                            ipset.port,
                            proto,
                            policy=FirewallRules.LOG_BLOCK_PORT_CHAIN,
                        ),
                        chain=FirewallRules.BP_INPUT_CHAIN,
                        table=FirewallRules.FILTER,
                        priority=FirewallRules.DEFAULT_PRIORITY,
                    )
                    for proto in self.expand_proto(ipset.proto)
                ]
            )
        return result

async def restore(self, ip_version: IPVersion) -> None:
        pass

async def get_ipsets_count(self, ip_version: IPVersion) -> list:
        return []

class IPSetIgnoredByPort(AbstractIPSet):
    async def block(self, ip, port, proto, *args, **kwargs):
        ipset = SingleIPSetPort(port=port, proto=proto)
        ip_version = IP.type_of(ip)
        await libipset.add_item(
            ipset.gen_ipset_name_for_ip_version(ip_version), ip, timeout=0
        )

async def unblock(self, ip, port, proto, *args, **kwargs):
        ipset = SingleIPSetPort(port=port, proto=proto)
        ip_version = IP.type_of(ip)
        await libipset.delete_item(
            ipset.gen_ipset_name_for_ip_version(ip_version), ip
        )

def get_all_ipset_instances(
        self, ip_version: IPVersion
    ) -> List[SingleIPSetPort]:
        if Firewall.port_blocking_mode != PORT_BLOCKING_MODE_ALLOW:
            return []
        result = []
        for port, proto in self._fetch(ip_version):
            result.append(SingleIPSetPort(port=port, proto=proto))
        return result

async def gen_ipset_restore_ops(self, ip_version: IPVersion):
        """
        Generate list of commands to fill all ip sets

:return: list of ipset commands to use with ipset restore
        """
        lines = []
        for ipset in self.get_all_ipset_instances(ip_version):
            lines.extend(await ipset.gen_ipset_restore_ops(ip_version))
        return lines

def _fetch(self, ip_version: IPVersion):
        return [
            (row.port_proto.port, row.port_proto.proto)
            for row in IgnoredByPort.fetch(
                NumericIPVersion.from_ip_version(ip_version)
            )
        ]

def gen_ipset_create_ops(self, ip_version: IPVersion) -> List[str]:
        # Just to follow an interface, there are no ipsets for this entity
        return []

def get_all_ipsets(self, ip_version: IPVersion) -> FrozenSet[str]:
        return frozenset()

def get_rules(self, ip_version: IPVersion, **kwargs) -> Iterable[dict]:
        return []

def is_enabled(self):
        return Firewall.port_blocking_mode == PORT_BLOCKING_MODE_ALLOW

async def restore(self, ip_version: IPVersion) -> None:
        if not self.is_enabled():
            return

with timeit("ipset_restore", logger):
            await libipset.restore(
                await self.gen_ipset_restore_ops(ip_version)
            )

async def get_ipsets_count(self, ip_version: IPVersion) -> list:
        if not self.is_enabled():
            return []

ipsets_count: Dict[str, int] = {}  # ipset name -> ips in db count
        for ipset in self.get_all_ipset_instances(ip_version):
            set_name = ipset.gen_ipset_name_for_ip_version(ip_version)
            ipsets_count[set_name] = ipset.db_count(ip_version)

ipsets = []
        for set_name, expected_count in ipsets_count.items():
            ipset_count = await libipset.get_ipset_count(set_name)
            ipsets.append(
                IPSetCount(
                    name=set_name,
                    db_count=expected_count,
                    ipset_count=ipset_count,
                )
            )
        return ipsets