ok
Direktori : /opt/imunify360/venv/lib/python3.11/site-packages/im360/internals/core/ipset/ |
Current File : //opt/imunify360/venv/lib/python3.11/site-packages/im360/internals/core/ipset/base.py |
import abc import hashlib import logging from functools import wraps from typing import List, Self from defence360agent.utils import LazyLock, log_error_and_ignore, timeit from im360.internals.core import ip_versions from im360.internals.core.ipset import libipset from im360.utils.validate import IPVersion logger = logging.getLogger(__name__) def ignore_if_ipset_not_found(coro): @wraps(coro) async def wrapper(self, *args, **kwargs): try: return await coro(self, *args, **kwargs) except libipset.IPSetNotFoundError: logger.warning( "%s not found. Skip '%s' command. Expect that actual ipset " "content will be restored from persistent during the next " "check of iptables rules/ipsets.", self.__class__.__qualname__, coro.__name__, ) return wrapper def raise_error_if_disabled(coro): @wraps(coro) async def wrapper(self, *args, **kwargs): if not self.is_enabled(): raise RuntimeError( f"Trying to use {coro.__name__} when " f"{self.__class__.__qualname__} is disabled." ) return await coro(self, *args, **kwargs) return wrapper class SingleIPSetInterface(abc.ABC): """ Interface for managing an IPSet. Implementing classes should represent a single IPSet. AbstractIPSet should be used as a collection of SingleIPSetInterface implementations. Example: - InputPortBlockingDenyModeIPSet: A collection of IP sets related to port blocking in deny mode (derived from AbstractIPSet). - i360.ipv6.input-ports-tcp, i360.ipv6.input-ports-udp: Instances of PortBlockingDenyModeIPSetManager (an implementation of SingleIPSetInterface). """ @abc.abstractmethod def gen_ipset_create_ops(self, ip_version: IPVersion) -> List[str]: pass @abc.abstractmethod def gen_ipset_destroy_ops(self, ip_version: IPVersion) -> List[str]: pass @abc.abstractmethod def gen_ipset_flush_ops(self, ip_version: IPVersion) -> List[str]: pass @abc.abstractmethod async def gen_ipset_restore_ops(self, ip_version: IPVersion) -> List[str]: pass @abc.abstractmethod def gen_ipset_name_for_ip_version(self, ip_version: IPVersion) -> str: pass def is_enabled(self, ip_version: IPVersion = None): return True @log_error_and_ignore( exception=libipset.IgnoredIPSetKernelError, log_handler=logger.warning ) @ignore_if_ipset_not_found @raise_error_if_disabled async def restore_from_persistent(self, ip_version: IPVersion): with timeit( "Restoring records for %s" % self.__class__.__name__, logger ): await libipset.restore( await self.gen_ipset_restore_ops(ip_version), name=self.gen_ipset_name_for_ip_version(ip_version), ) class IPSetAtomicRestoreBase(SingleIPSetInterface, abc.ABC): def __init__(self, *args, **kwargs): self.args = args self.kwargs = kwargs self.custom_ipset_name = None def clone_instance(self, new_ipset_name: str) -> Self: """ Create a copy of this instance and set a custom name for the related ipset. """ inst = self.__class__(*self.args, **self.kwargs) inst.custom_ipset_name = new_ipset_name return inst async def reset(self, ip_version: IPVersion = None): if ip_version: ip_versions_to_reset = [ip_version] else: ip_versions_to_reset = ip_versions.enabled() for ip_version in ip_versions_to_reset: logger.info( "Resetting %s ipset", self.gen_ipset_name_for_ip_version(ip_version), ) await self.restore_from_persistent_atomic(ip_version) async def exists(self, ip_version: IPVersion = None): name = self.gen_ipset_name_for_ip_version(ip_version) return name in await libipset.list_set() @staticmethod def get_tmp_ipset_name(original_ipset_name: str) -> str: tmp_ipset_hash = hashlib.sha1(original_ipset_name.encode()).hexdigest() return "i360." + tmp_ipset_hash[:20] + ".tmp" async def restore_from_persistent_atomic(self, ip_version: IPVersion): target_ipset_name = self.gen_ipset_name_for_ip_version(ip_version) tmp_ipset_name = self.get_tmp_ipset_name(target_ipset_name) tmp_ipset: IPSetAtomicRestoreBase = self.clone_instance(tmp_ipset_name) create_ipset_cmd = tmp_ipset.gen_ipset_create_ops(ip_version) if not create_ipset_cmd: return await libipset.restore(create_ipset_cmd) try: await libipset.flush_set(tmp_ipset_name) await tmp_ipset.restore_from_persistent(ip_version) await libipset.swap(tmp_ipset_name, target_ipset_name) finally: flush_ipset_cmd = tmp_ipset.gen_ipset_flush_ops(ip_version) destroy_ipset_cmd = tmp_ipset.gen_ipset_destroy_ops(ip_version) if flush_ipset_cmd: await libipset.restore(flush_ipset_cmd) if destroy_ipset_cmd: await libipset.restore(destroy_ipset_cmd)