1#@!#!123s

: / opt / imunify360 / venv / lib / python3.11 / site-packages / im360 / utils / tree_cache /
Filename : core.py
back
import itertools
import time
from abc import ABCMeta
from ipaddress import (
    IPV4LENGTH,
    IPV6LENGTH,
    IPv4Address,
    IPv4Network,
    IPv6Address,
    IPv6Network,
    ip_network,
)
from typing import Iterator, List, Optional, Union

import pytricia
from blinker import Signal

from im360.utils.validate import IP


class SourceInterface(metaclass=ABCMeta):
    # send **kwargs: ip, expiration
    added = Signal()

    # send **kwargs: ip
    deleted = Signal()

    # nothing is sent here
    cleared = Signal()

    # send **kwargs: ip, expiration
    updated = Signal()

    async def fetch_all(self):
        """:rtype: iterable[(ip, expiration)]"""
        raise NotImplementedError()


class TreeCacheInterface(metaclass=ABCMeta):
    async def contains(self, ip):
        """Check if the cache contains specified ip.

        :type ip: str
        """
        raise NotImplementedError()

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

        :type ip: str
        """
        raise NotImplementedError()

    async def filter_contained(
        self,
        ips: List[Union[IPv4Address, IPv4Network, IPv6Address, IPv6Network]],
    ):
        """Returns ips that is presented in cache."""
        raise NotImplementedError()

    async def filter_not_contained(
        self,
        ips: List[Union[IPv4Address, IPv4Network, IPv6Address, IPv6Network]],
    ):
        """Returns ips that is NOT presented in cache."""
        raise NotImplementedError()

    def reset(self):
        """
        re-fill on the next call
        """
        raise NotImplementedError()


class TreeCache(TreeCacheInterface):
    def __init__(self, source, full_sync_period=3600):
        """:type source: SourceInterface"""
        self._source = source
        self._full_sync_period = full_sync_period
        self._last_sync = 0

        # None means that we need initialize it on first call
        self._tree = None  # type: pytricia.PyTricia
        self._expired = False

        # subscribe to events from data source
        source.added.connect(self._on_added, source)
        source.deleted.connect(self._on_deleted, source)
        source.cleared.connect(self._on_cleared, source)
        source.updated.connect(self._on_updated, source)

    def _on_added(self, sender, ip, expiration):
        if self._tree is None:
            return
        self._tree.insert(IP.adopt_to_ipvX_network(ip), expiration)

    def _on_deleted(self, sender, ip):
        if self._tree is None:
            return
        try:
            # it's super fast even for very large trees
            self._tree.delete(IP.adopt_to_ipvX_network(ip))
        except KeyError:  # ip is not in tree
            pass

    def _on_cleared(self, sender):
        self._tree = pytricia.PyTricia(IPV6LENGTH)

    def _on_updated(self, sender, ip, expiration):
        self._on_added(sender, ip, expiration)

    async def _init_tree(self):
        self._tree = pytricia.PyTricia(IPV6LENGTH)
        tree = self._tree
        for ip, expiration in await self._source.fetch_all():
            tree.insert(IP.adopt_to_ipvX_network(ip), expiration)

    async def _sync_if_needed(self):
        while self._tree is None or time.time() > (
            self._last_sync + self._full_sync_period
        ):
            await self._init_tree()
            self._last_sync = time.time()

    @classmethod
    def _contains(
        cls,
        tree,
        ip: Union[IPv4Address, IPv4Network, IPv6Address, IPv6Network],
    ):
        """
        :return bool: if 'tree' contains ip as is or by subnet mask
        """
        ip_nwk, _ = cls._lookup(tree, ip)
        return bool(ip_nwk)

    @staticmethod
    def _lookup(
        tree, ip: Union[IPv4Address, IPv4Network, IPv6Address, IPv6Network]
    ):
        """Lookup specified ip or parent subnet considering expiration.
        :return tuple: str(ip_network(contains that ip)), expiration
        """
        ip = IP.adopt_to_ipvX_network(ip)
        # subnets can be nested, so check in loop.
        # endless loop is bad, so limit it.
        for _ in range(128):
            ip_nwk = tree.get_key(ip)

            if not ip_nwk:
                return None, None

            else:
                # check expiration
                expiration = tree.get(ip_nwk)
                if expiration and expiration <= time.time():
                    tree.delete(ip_nwk)
                    # check again
                    continue
                else:
                    return ip_nwk, expiration
        # For ipv4 max depth is 32, for ipv6 -- 128.
        # If all right, this place is unreachable.
        # If we are here, then something is wrong.
        raise RuntimeError(
            "Too deep recursion. Something goes wrong. "
            "Please contact developers."
        )

    @staticmethod
    def _contains_exactly(tree, ip):
        ip = IP.adopt_to_ipvX_network(ip)
        if tree.has_key(ip):  # noqa
            if TreeCache._contains(tree, ip):
                return tree.has_key(ip)  # noqa
        return False

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

        :type ip: str
        :return tuple: str(ip_network(contains that ip)), expiration
        """
        await self._sync_if_needed()
        return self._lookup(self._tree, ip)

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

        :type ip: str
        """
        await self._sync_if_needed()
        return self._contains(self._tree, ip)

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

        :type ip: str
        """
        await self._sync_if_needed()
        return self._contains_exactly(self._tree, ip)

    async def filter_contained(
        self,
        ips: List[Union[IPv4Address, IPv4Network, IPv6Address, IPv6Network]],
    ):
        """Returns ips that is presented in the cache."""
        if not ips:
            # a little optimization -- do not sync if ips is empty
            return ips

        await self._sync_if_needed()
        return [ip for ip in ips if self._contains(self._tree, ip)]

    async def filter_not_contained(
        self,
        ips: List[Union[IPv4Address, IPv4Network, IPv6Address, IPv6Network]],
    ):
        """Returns ips that is NOT presented in the cache."""
        if not ips:
            # a little optimization -- do not sync if ips is empty
            return ips

        await self._sync_if_needed()
        return [ip for ip in ips if not self._contains(self._tree, ip)]

    async def get_subnet(self, ip: str) -> Optional[str]:
        """Returns IP network that contains `ip`"""
        key = await self._get_key(ip)
        # as far as IPv6 addresses are stored in form of IPv6Network with
        # 64 bit mask, considering IP in the subnet only if key from cache
        # differs from original IP.
        # For example, "2002:add0:958a::/64" could be in tree cache and lookup
        # will return "2002:add0:958a::/64" too. Nevertheless, return value
        # is subnet, the function will return `False`, because `key != ip`
        # condition was unmet.
        if IP.is_valid_ip_network(key, strict=True) and key != ip:
            return key
        return None

    async def _get_key(self, ip_ntw):
        assert IP.is_valid_ip_network(ip_ntw), "%s is not valid IP!" % ip_ntw
        await self._sync_if_needed()
        return self._tree.get_key(ip_ntw)

    @staticmethod
    def _fix_addresses_in_network(network_list):
        # 10.1.1.1/32 -> 10.1.1.1
        return map(
            lambda ip: (
                str(ip_network(ip).network_address)
                if ip_network(ip).prefixlen == IPV4LENGTH
                else ip
            ),
            network_list,
        )

    async def get_ips_from_subnet(self, target_subnet: str) -> Iterator:
        """
        Returns a iterator of IP address and networks that located in cache
        and being members of `target_subnet` and `target_subnet` itself
        """
        if not IP.is_valid_ip_network(target_subnet, strict=True):
            return iter([])

        await self._sync_if_needed()
        parent_subnet = self._tree.get_key(target_subnet)
        if not parent_subnet:
            return iter([])

        items_in_subnet_from_cache = (
            net
            for net in self._tree.children(parent_subnet)
            if ip_network(net).network_address in ip_network(target_subnet)
        )
        return itertools.chain(
            self._fix_addresses_in_network(items_in_subnet_from_cache),
            [target_subnet],
        )

    def reset(self):
        """
        re-fill on the next call
        """
        # clear the state for _sync_if_needed() to init _trees
        self._tree = None