MAINHACK

Viewing / Editing File: ip.py

File is not writable. Editing disabled.
import dataclasses
import ipaddress
import itertools
import logging
import time
from abc import ABCMeta, abstractmethod
from typing import FrozenSet, Iterable, List, Self

from defence360agent.utils import log_error_and_ignore, timeit
from defence360agent.utils.common import DAY, rate_limit
from im360.api.server.ipecho import APIError as IPEchoAPIError
from im360.api.server.ipecho import IPEchoAPI
from im360.contracts.config import UnifiedAccessLogger
from im360.contracts.config import Webshield as WebshieldConfig
from im360.internals.core import rules
from im360.model.custom_lists import CustomBlacklist, CustomWhitelist
from im360.model.firewall import IgnoreList, IPList
from im360.model.firewall import IPv4 as DB_IPv4
from im360.model.firewall import IPv6 as DB_IPv6
from im360.model.firewall import RemoteProxy, RemoteProxyGroup
from im360.model.global_whitelist import (
    GlobalImunifyWhitelist,
    GlobalWhitelist,
)
from im360.subsys import webshield
from im360.utils.net import local_dns_from_resolv_conf, local_ip_addresses
from im360.utils.validate import IP, IPVersion, LocalhostIP

from .. import ip_versions
from ..firewall import FirewallRules, is_nat_available
from . import (
    IP_SET_PREFIX,
    AbstractIPSet,
    IPSetCount,
    get_ipset_family,
    libipset,
)
from .base import (
    IPSetAtomicRestoreBase,
    ignore_if_ipset_not_found,
    raise_error_if_disabled,
)
from .libipset import IPSetCmdBuilder
from .sync import (
    IPSetSyncCaptcha,
    IPSetSyncDrop,
    IPSetSyncSplashscreen,
    IPSetSyncWhite,
)

logger = logging.getLogger(__name__)
throttled_log_error = rate_limit(period=DAY, on_drop=logger.warning)(
    logger.error
)

ADD = 'add'
DEL = 'del'

_LOCAL_HOST_ADDRESS = {
    IP.V4: (4, LocalhostIP[IP.V4].value),
    IP.V6: (6, LocalhostIP[IP.V6].value),
}


def _prepare_command(ip, ipset_name, expiration=None, action=ADD,
                     ip_version=None):
    if (ip_version or IP.type_of(ip)) not in ip_versions.enabled():
        return None
    timeout = 0  # permanently
    if action == ADD:
        if expiration:
            timeout = int(expiration - time.time())
            if timeout <= 0:  # expired
                return None
    return " ".join(
        libipset.prepare_ipset_command(action, ipset_name, ip, timeout)) + "\n"


_RULE = dict(table=FirewallRules.FILTER,
             chain=FirewallRules.IMUNIFY_INPUT_CHAIN,
             priority=FirewallRules.DEFAULT_PRIORITY)


class BaseIPSet(IPSetAtomicRestoreBase, metaclass=ABCMeta):
    _NAME = ""  #: ipset name template such as '{prefix}.{ip_version}.graylist'
    DB_NAME = ""
    MAX_ELEM = 100000
    # according to
    # http://git.netfilter.org/ipset/tree/lib/parse.c#n1396
    # http://git.netfilter.org/ipset/tree/include/libipset/linux_ip_set.h
    MAX_SET_NAME_LENGTH = 31

    @log_error_and_ignore(
        exception=libipset.IgnoredIPSetKernelError, log_handler=logger.warning
    )
    @ignore_if_ipset_not_found
    @raise_error_if_disabled
    async def add(self, ip, timeout=0):
        version = IP.type_of(ip)
        if version not in ip_versions.enabled():
            logger.warning("Cannot add ip %s: %s is disabled", ip, version)
            return
        ipset_name = self.gen_ipset_name_for_ip_version(version)
        await libipset.add_item(ipset_name, ip, timeout)

    @log_error_and_ignore(
        exception=libipset.IgnoredIPSetKernelError, log_handler=logger.warning
    )
    @ignore_if_ipset_not_found
    @raise_error_if_disabled
    async def delete(self, ip):
        set_name = self._ipset_name_from_ip(ip)
        await libipset.delete_item(set_name, ip)

    async def get_db_count(self, ip_version: IPVersion):
        assert self.DB_NAME, "db name for set is not defined"
        iplist_version = (
            IPList.VERSION_IP4 if ip_version == IP.V4 else IPList.VERSION_IP6
        )
        return IPList.fetch_non_expired_query(
            self.DB_NAME, version=iplist_version).count()

    def _query(self, version):
        assert self.DB_NAME, "db name for set is not defined"
        return IPList.fetch_non_expired(self.DB_NAME, version=version)

    async def gen_ipset_restore_ops(
        self, ip_version: IPVersion
    ) -> List[str]:
        def prepare_command(_row, ipset_name):
            return _prepare_command(_row['ip'], ipset_name,
                                    _row.get('expiration'))

        ipset_name = self.gen_ipset_name_for_ip_version(ip_version)
        return list(
            filter(
                None,
                (
                    prepare_command(row, ipset_name)
                    for row in self._query(
                        version=(
                            IPList.VERSION_IP4
                            if ip_version == IP.V4
                            else IPList.VERSION_IP6
                        )
                    )
                ),
            )
        )

    @staticmethod
    def _make_record(ip, ipset_name, expiration=None, action=ADD) -> dict:
        return dict(
            ip=ip,
            ipset_name=ipset_name,
            expiration=expiration,
            action=action
        )

    @abstractmethod
    def rules(self, set_name: str, ip_version: IPVersion, **kwargs) \
            -> Iterable[dict]:  # pragma: no cover
        raise NotImplementedError()

    def gen_ipset_create_ops(
        self,
        ip_version: IPVersion,
        timeout: int = 0,
        datatype: str = libipset.HASH_NET,
        **options,
    ) -> List[str]:
        return [
            IPSetCmdBuilder.get_create_cmd(
                self.gen_ipset_name_for_ip_version(ip_version),
                datatype=datatype,
                family=get_ipset_family(ip_version),
                timeout=timeout,
                maxelem=self.MAX_ELEM,
            )
        ]

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

    def gen_ipset_name_for_ip_version(self, ip_version: IPVersion) -> str:
        if self.custom_ipset_name:
            return self.custom_ipset_name

        assert self._NAME, "set name is not defined"
        assert ip_version in (IP.V4, IP.V6), "IP {} is incorrect".format(
            ip_version
        )
        full_name = self._NAME.format(prefix=IP_SET_PREFIX,
                                      ip_version=ip_version)
        assert len(full_name) <= self.MAX_SET_NAME_LENGTH, \
            "setname {} is longer than {} characters".format(
                full_name, self.MAX_SET_NAME_LENGTH)
        return full_name

    def _ipset_name_from_ip(self, ip):
        return self.gen_ipset_name_for_ip_version(IP.type_of(ip))

    def create_rules(self, ip_version: IPVersion):
        set_name = self.gen_ipset_name_for_ip_version(ip_version)
        return self.rules(set_name, ip_version=ip_version)

    def _generate_for_restore(self, block_ips, unblock_ips):
        # first unblock then block
        for ip in unblock_ips:
            ipset_name = self._ipset_name_from_ip(ip)
            yield self._make_record(ip, ipset_name, action=DEL)

        for ip, properties in block_ips:
            ipset_name = self._ipset_name_from_ip(ip)
            yield self._make_record(
                ip=ip, ipset_name=ipset_name,
                expiration=properties['expiration']
            )

    @log_error_and_ignore(
        exception=libipset.IgnoredIPSetKernelError, log_handler=logger.warning
    )
    @ignore_if_ipset_not_found
    @raise_error_if_disabled
    async def restore(self, block_ips, unblock_ips):
        """Run "ipset restore" for given ips."""
        lines = list()
        records = self._generate_for_restore(block_ips, unblock_ips)
        for rec in records:
            cmd = _prepare_command(**rec)
            if cmd:
                lines.append(cmd)
        with timeit(
            "ipset_restore [%s]" % (self.__class__.__name__, ),
            logger
        ):
            await libipset.restore(lines, name=self.__class__.__name__)

    def _log_rule(self, set_name, ip_version: IPVersion, prefix, priority):
        yield from map(
            dataclasses.asdict,
            rules.log_rules(set_name, ip_version, prefix, priority),
        )


class WebshieldEnabledIPSet(BaseIPSet):
    def is_enabled(self, ip_version: IPVersion = None) -> bool:
        if not super().is_enabled():
            return False  # short circuit behavior

        enabled = WebshieldConfig.ENABLE
        if enabled and not (
            webshield_expects_traffic := webshield.expects_traffic()
        ):
            throttled_log_error(
                "Webshield enabled, but it does not expect traffic"
            )
        return enabled and webshield_expects_traffic


class IPSetGray(WebshieldEnabledIPSet):
    _NAME = '{prefix}.{ip_version}.graylist'
    DB_NAME = IPList.GRAY
    MAX_ELEM = 2000000
    # Default block in the graylist ipset in days
    GRAYLIST_DEFAULT_TIMEOUT = libipset.IPSET_TIMEOUT_MAX

    def rules(
        self, set_name: str, ip_version: IPVersion, **kwargs
    ) -> Iterable[dict]:
        yield from map(
            dataclasses.asdict,
            rules.webshield_rules(
                set_name,
                ip_version,
                # note: make the flag is true in exactly one place
                rules.CaptchaRuleBuilder(),
            ),
        )

    def gen_ipset_create_ops(
        self,
        ip_version: IPVersion,
        timeout: int = 0,
        datatype: str = libipset.HASH_NET,
        **options,
    ) -> List[str]:
        return super().gen_ipset_create_ops(
            ip_version=ip_version,
            timeout=self.GRAYLIST_DEFAULT_TIMEOUT,
        )


class IPSetGraySplashScreen(IPSetGray):
    """ Inherited from Gray this list has less priority and do not block,
    only redirect to webshield webports. """

    _NAME = "{prefix}.{ip_version}.graysplashlist"
    DB_NAME = IPList.GRAY_SPLASHSCREEN

    def is_enabled(self, ip_version: IPVersion = None) -> bool:
        return super().is_enabled() and WebshieldConfig.SPLASH_SCREEN

    def rules(
        self, set_name: str, ip_version: IPVersion, **kwargs
    ) -> Iterable[dict]:
        yield from map(
            dataclasses.asdict,
            rules.webshield_rules(
                set_name, ip_version, rules.SplashscreenRuleBuilder()
            ),
        )


class IPSetRemoteProxy(WebshieldEnabledIPSet):
    _NAME = '{prefix}.{ip_version}.remote_proxy'
    DB_NAME = ''  # we override _query

    def rules(self, set_name: str, ip_version: IPVersion, **kwargs) \
            -> Iterable[dict]:
        redirect_map = webshield.port_redirect_map()
        dest_ports = webshield.redirected_to_webshield_ports() & set(
            redirect_map
        )
        yield from map(
            dataclasses.asdict,
            rules.check_access_to_webshield_ports_rules(
                set_name, set(redirect_map[p] for p in dest_ports)
            ),
        )

        yield dict(_RULE,
                   rule=FirewallRules.open_dst_ports_for_src_list(
                       set_name, set(redirect_map[p] for p in dest_ports)),
                   priority=FirewallRules.REMOTE_PROXY_PRIORITY)

        if is_nat_available(ip_version):
            yield from map(
                dataclasses.asdict,
                rules.redirect_port_rules(
                    set_name,
                    dest_ports,
                    redirect_map,
                    FirewallRules.NAT,
                    FirewallRules.redirect_to_captcha,
                ),
            )
        else:
            # Similar to IPSetGray
            yield from map(
                dataclasses.asdict,
                rules.redirect_port_rules(
                    set_name,
                    dest_ports,
                    redirect_map,
                    FirewallRules.MANGLE,
                    FirewallRules.redirect_to_captcha_via_tproxy,
                ),
            )

            yield dict(
                _RULE, rule=FirewallRules.traffic_not_from_tproxy(set_name)
            )

    async def get_db_count(self, ip_version: IPVersion):
        iplist_version = (
            IPList.VERSION_IP4 if ip_version == IP.V4 else IPList.VERSION_IP6
        )
        q = (
            RemoteProxy.select(RemoteProxy.network)
            .join(RemoteProxyGroup)
            .where(RemoteProxyGroup.enabled)
        )
        return sum(
            ipaddress.ip_network(item[0]).version == iplist_version
            for item in q.tuples()
        )

    def _query(self, version):
        q = (
            RemoteProxy.select(RemoteProxy.network)
            .join(RemoteProxyGroup)
            .where(RemoteProxyGroup.enabled)
        )
        return [
            {"ip": item[0], "expiration": 0}
            for item in q.tuples()
            if ipaddress.ip_network(item[0]).version == version
        ]


class IPSetStaticRemoteProxy(IPSetRemoteProxy):
    _NAME = '{prefix}.{ip_version}.remote_proxy_static'

    async def gen_ipset_restore_ops(self, ip_version: IPVersion) -> List[str]:
        cmd_list = []
        ipset_name = self.gen_ipset_name_for_ip_version(ip_version)
        for ip in await GlobalWhitelist.load(group='proxy'):
            if IP.type_of(ip) != ip_version:
                continue
            cmd = _prepare_command(ip, ipset_name, ip_version=ip_version)
            if cmd:
                cmd_list.append(cmd)
        return cmd_list

    def is_enabled(self, ip_version: IPVersion = None) -> bool:
        return super().is_enabled() and WebshieldConfig.KNOWN_PROXIES_SUPPORT

    async def get_db_count(self, ip_version: IPVersion):
        return sum(
            IP.type_of(ip) == ip_version
            for ip in await GlobalWhitelist.load(group="proxy")
        )


class IPSetWhite(BaseIPSet):
    _NAME = '{prefix}.{ip_version}.whitelist'
    DB_NAME = IPList.WHITE

    def rules(self, set_name: str, ip_version: IPVersion, **kwargs) \
            -> Iterable[dict]:
        yield from self._log_rule(
            set_name, ip_version,
            UnifiedAccessLogger.WHITELIST, FirewallRules.WHITELIST_PRIORITY)
        yield dict(_RULE,
                   rule=FirewallRules.ipset_rule(set_name,
                                                 FirewallRules.RETURN),
                   priority=FirewallRules.WHITELIST_PRIORITY)

        if is_nat_available(ip_version):
            yield dict(_RULE,
                       rule=FirewallRules.ipset_rule(set_name,
                                                     FirewallRules.RETURN),
                       table=FirewallRules.NAT)
        else:
            yield dict(_RULE,
                       rule=FirewallRules.ipset_rule(set_name,
                                                     FirewallRules.RETURN),
                       table=FirewallRules.MANGLE)

    async def delete(self, ip):
        await super(IPSetWhite, self).delete(ip)

        # we need also delete from full access list
        await IPSetWhiteFullAccess().delete(ip)

    async def get_db_count(self, ip_version: IPVersion):
        db_ip_version = DB_IPv4 if ip_version == IP.V4 else DB_IPv6
        return IPList.fetch_non_expired_query(
            IPList.WHITE, full_access=False, version=db_ip_version).count()

    def _query(self, version):
        return IPList.fetch_non_expired(IPList.WHITE,
                                        full_access=False,
                                        version=version)

    def get_non_captcha_passed_ips(self):
        whitelisted_entries = IPList.fetch_non_expired_query(
            IPList.WHITE, full_access=False)
        non_captcha_passed_entries = whitelisted_entries.where(
            ~IPList.captcha_passed).dicts().iterator()
        # FIXME: after migrating to peewee 3 switch back to this
        # return (entry["ip"] for entry in non_captcha_passed_entries)
        try:
            for entry in non_captcha_passed_entries:
                yield entry["ip"]
        except RuntimeError:
            return


class IPSetBlack(BaseIPSet):
    _NAME = '{prefix}.{ip_version}.blacklist'
    DB_NAME = IPList.BLACK

    def rules(self, set_name: str, ip_version: IPVersion, **kwargs) \
            -> Iterable[dict]:
        yield from map(
            dataclasses.asdict, rules.drop_rules(set_name, ip_version)
        )


class IPSetWhiteFullAccess(BaseIPSet):
    _NAME = '{prefix}.{ip_version}.whitelist.full_access'

    def rules(self, set_name: str, ip_version: IPVersion, **kwargs) \
            -> Iterable[dict]:
        yield from self._log_rule(
            set_name, ip_version,
            UnifiedAccessLogger.WHITELIST,
            FirewallRules.FULL_ACCESS_PRIORITY)
        yield dict(_RULE,
                   rule=FirewallRules.ipset_rule(set_name,
                                                 FirewallRules.RETURN),
                   priority=FirewallRules.FULL_ACCESS_PRIORITY)
        yield dict(
            table=FirewallRules.FILTER,
            chain=FirewallRules.IMUNIFY_OUTPUT_CHAIN,
            priority=FirewallRules.FULL_ACCESS_PRIORITY,
            # it is much better to write iptables rules explicitly, instead
            # of guessing that somewhere in the deepest stack frames it uses
            # 'src' instead of desired 'dst'
            rule=('-m', 'set', '--match-set',
                  set_name, 'dst', '-j', FirewallRules.RETURN),
        )
        if is_nat_available(ip_version):
            yield dict(_RULE,
                       rule=FirewallRules.ipset_rule(set_name,
                                                     FirewallRules.RETURN),
                       table=FirewallRules.NAT)
        else:
            yield dict(_RULE,
                       rule=FirewallRules.ipset_rule(set_name,
                                                     FirewallRules.RETURN),
                       table=FirewallRules.MANGLE)

    def get_local_records(self, version=None):
        for ip_version, (_version, ip_address) in _LOCAL_HOST_ADDRESS.items():
            if not version or _version == version:
                yield self._make_record(
                    ip_address, self.gen_ipset_name_for_ip_version(ip_version)
                )

    async def get_db_count(self, ip_version: IPVersion):
        db_ip_version = DB_IPv4 if ip_version == IP.V4 else DB_IPv6
        local_count = sum(
            1 for _ in self.get_local_records(version=db_ip_version)
        )
        whitelisted_db_count = IPList.fetch_non_expired_query(
            IPList.WHITE, full_access=True, version=db_ip_version).count()
        return local_count + whitelisted_db_count

    def _query(self, version=None):
        return itertools.chain(
            self.get_local_records(version=version),
            IPList.fetch_non_expired(IPList.WHITE,
                                     full_access=True,
                                     version=version))

    def query_all(self):
        return self._query()


class IPSetStatic(BaseIPSet):
    _NAME = '{prefix}.{ip_version}.whitelist.static'
    _PRIORITY = FirewallRules.STATIC_WHITELIST_PRIORITY

    def rules(self, set_name: str, ip_version: IPVersion, **kwargs) \
            -> Iterable[dict]:
        yield from map(
            dataclasses.asdict,
            rules.white_rules(
                set_name,
                ip_version,
                priority=self._PRIORITY,
            ),
        )

    async def gen_ipset_restore_ops(self, ip_version: IPVersion):
        cmd_list = []
        ipset_name = self.gen_ipset_name_for_ip_version(ip_version)
        for ip in await self._get_ips(ip_version):
            if IP.type_of(ip) != ip_version:
                continue
            cmd = _prepare_command(ip, ipset_name, ip_version=ip_version)
            if cmd:
                cmd_list.append(cmd)
        return cmd_list

    async def _get_ips(self, ip_version: IPVersion):
        return await GlobalWhitelist.load()

    async def get_db_count(self, ip_version: IPVersion):
        return sum(
            IP.type_of(ip) == ip_version
            for ip in await self._get_ips(ip_version)
        )


class IPSetI360Static(IPSetStatic):
    _NAME = "{prefix}.{ip_version}.i360_whitelist.static"
    _PRIORITY = FirewallRules.WHITELIST_PRIORITY

    async def _get_ips(self, ip_version: IPVersion):
        return await GlobalImunifyWhitelist.load()


class IPSetWhitelistHostIPs(IPSetStatic):
    _NAME = '{prefix}.{ip_version}.whitelist.host_ips'
    _PRIORITY = FirewallRules.HOST_IPS_PRIORITY

    async def _get_ips(self, ip_version: IPVersion):
        local_ips = map(str, map(IP.ipv6_to_64network, local_ip_addresses()))
        try:
            own_nat_ip = await IPEchoAPI.get_ip(ip_version)
        except IPEchoAPIError:
            own_nat_ip = None
        return itertools.chain(
            local_dns_from_resolv_conf(),
            local_ips,
            [str(IP.ipv6_to_64network(own_nat_ip))] if own_nat_ip else []
        )


class IPSetCustomWhitelist(IPSetStatic):
    _NAME = '{prefix}.{ip_version}.whitelist.custom'
    _LIST = CustomWhitelist
    _PRIORITY = FirewallRules.WHITELIST_PRIORITY
    MAX_ELEM = 524288

    async def gen_ipset_restore_ops(self, ip_version: IPVersion) -> List[str]:
        cmd_list = []
        ipset_name = self.gen_ipset_name_for_ip_version(ip_version)
        for ip in await self._LIST.load():
            if IP.type_of(ip) != ip_version:
                continue
            cmd = _prepare_command(ip, ipset_name, ip_version=ip_version)
            if cmd:
                cmd_list.append(cmd)
        return cmd_list

    async def get_db_count(self, ip_version: IPVersion):
        return sum(
            IP.type_of(ip) == ip_version for ip in await self._LIST.load()
        )


class IPSetCustomBlacklist(IPSetCustomWhitelist):
    _NAME = '{prefix}.{ip_version}.blacklist.custom'
    _LIST = CustomBlacklist

    def rules(self, set_name: str, ip_version: IPVersion, **kwargs) \
            -> Iterable[dict]:
        yield dict(_RULE,
                   rule=FirewallRules.ipset_rule(
                       set_name, FirewallRules.LOG_BLACKLIST_CHAIN),
                   priority=FirewallRules.BLACKLIST_PRIORITY)


class IPSetIgnore(BaseIPSet):
    _NAME = '{prefix}.{ip_version}.ignorelist'

    def rules(self, set_name: str, ip_version: IPVersion, **kwargs) \
            -> Iterable[dict]:
        yield dict(_RULE,
                   rule=FirewallRules.ipset_rule(set_name,
                                                 FirewallRules.RETURN))
        if is_nat_available(ip_version):
            yield dict(_RULE,
                       rule=FirewallRules.ipset_rule(set_name,
                                                     FirewallRules.RETURN),
                       table=FirewallRules.NAT)
        else:
            yield dict(
                _RULE,
                rule=FirewallRules.ipset_rule(set_name, FirewallRules.RETURN),
                table=FirewallRules.MANGLE,
            )

    async def get_db_count(self, ip_version: IPVersion):
        db_ip_version = DB_IPv4 if ip_version == IP.V4 else DB_IPv6
        return (
            IgnoreList.select()
            .where(IgnoreList.version == db_ip_version)
            .count()
        )

    def _query(self, version):
        return IgnoreList.select().where(IgnoreList.version == version).dicts()


class IPSet(AbstractIPSet):
    def __init__(self):
        super().__init__()
        self.ip_sets = [
            IPSetRemoteProxy(),
            IPSetStaticRemoteProxy(),
            IPSetWhiteFullAccess(),
            IPSetStatic(),
            IPSetI360Static(),
            IPSetWhitelistHostIPs(),
            IPSetSyncWhite(),
            IPSetCustomWhitelist(),
            IPSetWhite(),
            IPSetIgnore(),
            IPSetBlack(),
            IPSetSyncDrop(),
            IPSetCustomBlacklist(),
            IPSetGraySplashScreen(),
            IPSetSyncSplashscreen(),
            IPSetGray(),
            IPSetSyncCaptcha(),
        ]

    def get_all_ipsets(self, ip_version: IPVersion) -> FrozenSet[str]:
        return frozenset(
            set_.gen_ipset_name_for_ip_version(ip_version)
            for set_ in self.ip_sets
            if set_.is_enabled()
        )

    def get_all_ipset_instances(
        self, ip_version: IPVersion
    ) -> List[IPSetAtomicRestoreBase]:
        return self.ip_sets

    def get_ipset(self, db_listname):
        for ipset_ in self.ip_sets:
            if ipset_.DB_NAME == db_listname:
                return ipset_
        raise LookupError("Set {} not found".format(db_listname))

    async def block(self, ip, listname=IPList.GRAY, timeout=0,
                    full_access=False, *args, **kwargs):
        """Block the ip

        :param ip: ip for blocking
        :param listname: ipset list for blocking
        :param timeout: relative timeout in seconds, if equal 0 - permanently
        :param full_access: full access for whitelist
        :return:
        """

        assert IP.is_valid_ip_network(ip)

        ipset_ = IPSetWhiteFullAccess() if full_access else self.get_ipset(
            listname)
        if not ipset_.is_enabled():
            raise RuntimeError("Set {} is disabled".format(
                ipset_.__class__.__name__))
        await ipset_.add(ip, timeout)

    async def unblock(self, ip, listname=IPList.GRAY, *args, **kwargs):
        """Unblock the ip

        :param ip: ip for blocking
        :param listname: ipset list for blocking
        :return:
        """
        assert IP.is_valid_ip_network(ip)
        set_ = self.get_ipset(listname)
        if set_.is_enabled():
            await set_.delete(ip)

    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
        """
        ipsets = []
        for set_ in self.ip_sets:
            if set_.is_enabled():
                ipsets.extend(set_.gen_ipset_create_ops(ip_version))
        return ipsets

    def get_rules(self, ip_version: IPVersion, **kwargs) -> Iterable[dict]:
        ruleset = []
        for set_ in self.ip_sets:
            if set_.is_enabled():
                ruleset.extend(set_.create_rules(ip_version))
        return ruleset

    async def restore(self, ip_version: IPVersion) -> None:
        for s in self.ip_sets:
            if s.is_enabled():
                await s.restore_from_persistent(ip_version)

    async def get_ipsets_count(self, ip_version: IPVersion) -> list:
        ipsets = []
        for ip_set in self.ip_sets:
            if ip_set.is_enabled():
                set_name = ip_set.gen_ipset_name_for_ip_version(ip_version)
                expected_count = await ip_set.get_db_count(ip_version)
                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
Cancel / Back
[ MAINHACK ]

Viewing / Editing File: ip.py

Server Info

System Features

User Info
