1#@!#!123s

: / opt / imunify360 / venv / lib / python3.11 / site-packages / im360 / utils / tree_cache /
Filename : adapters.py
back
from ipaddress import IPv4Address, IPv4Network, IPv6Address, IPv6Network
from typing import List, Union

from .core import TreeCacheInterface


class TreeCacheCollector(TreeCacheInterface):
    """
    Facade for multiple caches.
    :param collected_caches: pool of TreeCacheInterface
    """

    def __init__(self, *collected_caches):
        self._collected_caches = collected_caches

    async def contains(self, ip):
        """Check if the cache contains specified ip.

        :type ip: str
        """
        for cache in self._collected_caches:
            if await cache.contains(ip):
                return True
        return False

    async def contains_exactly(self, ip):
        """Check if the cache contains exactly specified ip, not parent subnet.

        :type ip: str
        """
        for cache in self._collected_caches:
            if await cache.contains_exactly(ip):
                return True
        return False

    async def filter_contained(
        self,
        ips: List[Union[IPv4Address, IPv4Network, IPv6Address, IPv6Network]],
    ):
        """Returns ips those are presented in cache."""
        set_ = set()
        for cache in self._collected_caches:
            for ip in await cache.filter_contained(ips):
                set_.add(ip)
        return list(set_)

    async def filter_not_contained(
        self,
        ips: List[Union[IPv4Address, IPv4Network, IPv6Address, IPv6Network]],
    ):
        """Returns ips those are NOT presented in cache."""
        rv = ips
        for cache in self._collected_caches:
            if rv:
                rv = await cache.filter_not_contained(rv)
            else:
                break
        return rv

    def reset(self):
        """
        Each of collected_caches will be re-init from so-called SourceInterface
        on the next call
        """
        for treecache_interface in self._collected_caches:
            treecache_interface.reset()