Current File : //opt/imunify360/venv/lib/python3.11/site-packages/im360/utils/tree_cache/core.py
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
Mini Shell