MAINHACK

Viewing / Editing File: lazy_init.py

File is not writable. Editing disabled.
"""This plugin periodically checks set of rules and ipsets,
and recreates it if needed, process block/unblock messages."""

import asyncio
import json
import logging
import os
import time
from contextlib import suppress
from typing import Set

from defence360agent import utils
from defence360agent.contracts.config import UserConfig
from defence360agent.contracts.messages import MessageType
from defence360agent.contracts.plugins import MessageSink, expect
from defence360agent.utils import log_error_and_ignore, recurring_check, timeit
from defence360agent.utils.common import DAY, ServiceBase, rate_limit
from im360.contracts.config import Firewall, NetworkInterface
from im360.contracts.config import Protector as Config
from im360.contracts.config import Webshield
from im360.files import GEO, WHITELISTS
from im360.internals.core import IPSetNoRedirectPort, RuleSet, ip_versions
from im360.internals.core.firewall import get_firewall
from im360.internals.core.firewall.base import (
    FirewallBatchCommandError,
    FirewallError,
    FirewallTemporaryError,
)
from im360.internals.core.ipset.country import IPSetCountry
from im360.internals.core.ipset.ip import (
    IPSetBlack,
    IPSetCustomBlacklist,
    IPSetCustomWhitelist,
    IPSetGray,
    IPSetGraySplashScreen,
    IPSetIgnore,
    IPSetStatic,
    IPSetStaticRemoteProxy,
    IPSetWhite,
)
from im360.internals.core.ipset.port_deny import (
    InputPortBlockingDenyModeIPSet,
    OutputPortBlockingDenyModeIPSet,
)
from im360.model.firewall import IPList
from im360.utils.validate import IP, IPVersion

logger = logging.getLogger(__name__)


class VersionState:
    def __init__(self):
        self.errors = 0
        self.next_try_time = 0.0
        self.running = True


class RulesChecker(ServiceBase):
    """Periodically checks if rules exist and if not, recreate them."""
    #: delay in secs between two iptables rule/ipsets checks
    RULE_CHECK_INTERVAL = int(
        os.environ.get("IMUNIFY360_RULE_CHECK_INTERVAL", 30)
    )
    CONSECUTIVE_ERROR_LIMIT = 10  # retries until giving up
    ERROR_THRESHOLD = 10  # disabled IPv6 support after this many failures
    RETRY_INTERVAL = 3600  # interval between attempts to enable IPv6
    CAPTURE_CSF_LOCK = True
    CSF_LOCK_TIMEOUT = 300  # seconds
    IPSETS_CHECK_INTERVAL = DAY  # seconds

    def __init__(self, loop):
        super().__init__(loop)
        self.ruleset = RuleSet()
        # Prevent rules create/destroy public method calls
        # from overlapping with @recurring_check->create_rules/destroy_rules()
        #
        # Use this lock to operate on rules (add to ipset)
        # while keeping the rules in consistent state (not half created
        # or half destroyed).
        self.lock_rules_create_destroy = Config.RULE_EDIT_LOCK

        self.active_interface_conf = NetworkInterface.get_interface_conf()
        self.versions = {ver: VersionState() for ver in ip_versions.all()}
        self._ipsets_outdated_events = {
            ver: asyncio.Event()
            for ver in self.versions
        }
        self.ipset_lock = asyncio.Lock()

    async def recreate_rules_if_needed(self, recreate_any_way=False) -> None:
        # re-reading config file
        async with self.lock_rules_create_destroy:
            new_conf = NetworkInterface.get_interface_conf()
            if not new_conf == self.active_interface_conf:
                logger.info('Target & ignore interfaces config was changed,'
                            'recreating rules')
                await self._destroy_rules_and_sets(self.active_interface_conf)
            await self._ensure_rules_exist_for_all_versions(new_conf,
                                                            recreate_any_way)
            self.active_interface_conf = new_conf

    @recurring_check(RULE_CHECK_INTERVAL, CONSECUTIVE_ERROR_LIMIT)
    async def _run(self):
        await self.recreate_rules_if_needed()
        await self.check_ipsets_consistent()  # pragma: no cover

    @rate_limit(period=IPSETS_CHECK_INTERVAL)
    @log_error_and_ignore()
    async def check_ipsets_consistent(self):  # pragma: no cover
        await self._check_ipsets_consistent()

    async def _check_ipsets_consistent(self):
        async with self.lock_rules_create_destroy:
            outdated_ipsets = []
            for version, state in self.versions.items():
                if state.running:
                    if (
                        outdated := await self.ruleset.get_outdated_ipsets(
                            version
                        )
                    ):
                        outdated_ipsets += outdated
                        self._ipsets_outdated_events[version].set()
            if outdated_ipsets:
                logger.warning(
                    "%s ipsets are outdated. Recreate them", outdated_ipsets
                )

    async def _ensure_rules(self, interface_conf,
                            recreate_any_way=False) -> bool:
        now = time.monotonic()
        target_versions = set(
            version
            for version in ip_versions.all()
            if (ip_versions.is_enabled(version)
                or self.versions[version].next_try_time < now))
        recreated_any = False
        has_failed = False
        for version in target_versions:
            self.versions[version].running = False
            try:
                async with self.ipset_lock:
                    r = await self._ensure_for(
                        version, interface_conf, recreate_any_way
                    )
                recreated_any = recreated_any or r
            except FirewallTemporaryError as exc:
                has_failed = True
                logger.info('Transient error while creating firewall rules for'
                            ' %s: %s', version, exc)
            except FirewallError as exc:
                enabled_ip_version = ip_versions.is_enabled(version)
                if enabled_ip_version:
                    has_failed = True
                    self.versions[version].errors += 1
                    if self.versions[version].errors >= self.ERROR_THRESHOLD:
                        if ip_versions.disable(version):
                            self.versions[version].next_try_time = (
                                now + self.RETRY_INTERVAL)
                            logger.info('%s firewall support is disabled due'
                                        ' to multiple consecutive errors',
                                        version)
                else:
                    self.versions[version].next_try_time = (
                        now + self.RETRY_INTERVAL)
                logger.warning(
                    'Failed to recreate firewall rules for %s %s: %s',
                    "enabled" if enabled_ip_version else "disabled",
                    version, exc)
            else:
                self.versions[version].running = True
                self.versions[version].errors = 0
                if not ip_versions.is_enabled(version):
                    ip_versions.enable(version)
                    logger.info('%s firewall support is enabled', version)
        # whether recreated rules for any ip version and no failures
        # in creating rules for enabled ip versions
        return recreated_any and not has_failed

    async def _ensure_rules_exist_for_all_versions(self, interface_conf=None,
                                                   recreate_any_way=False):
        interface_conf = interface_conf or self.active_interface_conf
        if await self._ensure_rules(interface_conf, recreate_any_way):
            logger.info('Rules and sets successfully '
                        'recreated for enabled ip versions')

    async def _recreate_for(
        self,
        ip_version: IPVersion,
        interface_conf: dict,
        recreate_ipsets: bool,
        existing_ipsets: Set[str],
        missing_ipsets: Set[str],
        redundant_ipsets: Set[str],
    ):
        with timeit(f"Destroying rules for {ip_version}", log=logger.info):
            await self._destroy_rules(interface_conf, ip_version)

        if redundant_ipsets:
            with timeit(
                f"Destroying redundant ipsets: {redundant_ipsets}",
                log=logger.info,
            ):
                await self.ruleset.destroy_ipsets(ip_version, redundant_ipsets)

        if missing_ipsets:
            with timeit(
                f"Creating missing ipsets: {missing_ipsets}", log=logger.info
            ):
                await self.ruleset.fill_ipsets(ip_version, missing_ipsets)

        with timeit(f"Recreating rules for {ip_version}", log=logger.info):
            await self._create_rules(interface_conf, ip_version)

        if recreate_ipsets:
            with timeit("Recreating ipsets", log=logger.info):
                await self.ruleset.recreate_ipsets(ip_version, existing_ipsets)

    async def _ensure_for(
        self,
        ip_version: IPVersion,
        interface_conf: dict,
        recreate_any_way=False,
    ) -> bool:
        """Creates imunify360 ruleset for given IP version in iptables.

        If all required ipsets, rules and chains exist, does nothing.
        Otherwise recreates everything as required.

        Returns True if rules or sets has been (re-)created, False
        otherwise."""
        existing_ipsets = await self.ruleset.existing_ipsets(ip_version)
        required_ipsets = self.ruleset.required_ipsets(ip_version)
        redundant_ipsets = set()
        missing_ipsets = required_ipsets
        ipsets_ok = (
            existing_ipsets == required_ipsets
            and not self._ipsets_outdated_events[ip_version].is_set()
        )
        # don't log in case when no existing ipsets, it is OK on start
        if existing_ipsets and existing_ipsets != required_ipsets:
            redundant_ipsets = existing_ipsets - required_ipsets
            missing_ipsets -= existing_ipsets
            existing_ipsets -= redundant_ipsets
            _log_ipsets_mismatch(
                missing_ipsets=missing_ipsets,
                redundant_ipsets=redundant_ipsets,
                log=logger.warning,
            )
        rules_ok = await self._rules_ok(interface_conf, ip_version)
        if not (rules_ok and ipsets_ok and not recreate_any_way):
            if not ipsets_ok:  # ipsets will be re-created
                self._ipsets_outdated_events[ip_version].clear()
            logger.info(
                "Rules status for %s [rules: %s], [ipset: %s]",
                ip_version,
                "ok" if rules_ok else "bad",
                "ok" if ipsets_ok else "bad")
            await self._recreate_for(
                ip_version,
                interface_conf,
                recreate_ipsets=not ipsets_ok,
                existing_ipsets=existing_ipsets,
                redundant_ipsets=redundant_ipsets,
                missing_ipsets=missing_ipsets,
            )
            return True
        return False

    async def _create_rules(
        self, interface_conf, ip_version: IPVersion
    ) -> None:
        async with await get_firewall(ip_version) as firewall:
            batch = await self.ruleset.create_commands(
                firewall, interface_conf, ip_version)
            await firewall.commit(batch)

    async def _destroy_rules(
        self, interface_conf, ip_version: IPVersion
    ) -> None:
        async with await get_firewall(ip_version) as firewall:
            for batch in self.ruleset.destroy_commands(
                    firewall, interface_conf, ip_version):
                with suppress(FirewallBatchCommandError):
                    # Command may fail and that is ok
                    await firewall.commit(batch)

    async def _rules_ok(self, interface_conf, ip_version: IPVersion) -> bool:
        async with await get_firewall(ip_version) as firewall:
            actions = await self.ruleset.check_commands(
                firewall, interface_conf, ip_version)
            try:
                await firewall.commit(actions)
            except FirewallBatchCommandError:
                return False
            else:
                return True

    async def _destroy_rules_and_sets(self, interface_conf):
        errors = []
        for ip_version in ip_versions.enabled():
            self.versions[ip_version].running = False
            try:
                await self._destroy_rules(interface_conf, ip_version)
                await self.ruleset.destroy_ipsets(ip_version)
            except Exception as e:
                errors.append(e)

        if not errors:
            return
        elif len(errors) == 1:
            raise errors[0]
        elif len(errors) == 2:
            # hack: to get "free" readable traceback for both errors
            raise errors[1] from errors[0]
        else:  # pragma: no cover
            assert 0, "max 2 ip versions expected"

    async def clear_everything(self, interface_conf=None) -> None:
        interface_conf = interface_conf or self.active_interface_conf
        async with self.lock_rules_create_destroy:
            await self._destroy_rules_and_sets(interface_conf)


class RealProtector(MessageSink):
    PROCESSING_ORDER = MessageSink.ProcessingOrder.IPSET_PROTECTOR

    def __init__(self):
        super().__init__()
        self._port_blocking_mode = Firewall.port_blocking_mode
        self._port_blocking_deny_mode_values \
            = self._get_port_blocking_deny_mode_values()

    async def create_sink(self, loop, rules_checker=None):
        """Create all ipset sets"""
        self._loop = loop
        self._process_config_update_lock = asyncio.Lock()
        self._shutdown_lock = asyncio.Lock()
        # saving webshield status before updating the rules
        # (on _rules_checker.start()), to avoid re-checking
        # the rules on the 1st ConfigUpdate (on start-up)
        self._webshield_status = (
            Webshield.ENABLE,
            Webshield.SPLASH_SCREEN,
            Webshield.PANEL_PROTECTION,
        )
        self._rules_checker = rules_checker or RulesChecker(loop)
        self._rules_checker.start()

    async def shutdown(self):
        async with self._shutdown_lock:
            self._rules_checker.should_stop()
            await self._rules_checker.wait()
        await self._rules_checker.clear_everything()
        logger.info("IP sets and rules have been removed successfully.")

    def _get_port_blocking_deny_mode_values(self):
        return (
            Firewall.TCP_IN_IPV4,
            Firewall.TCP_OUT_IPV4,
            Firewall.UDP_IN_IPV4,
            Firewall.UDP_OUT_IPV4,
        )

    def _get_current_ttl(self, expiration):
        if expiration:
            return int(expiration - time.time())
        return 0

    def _rules_running(self, ip_version: IPVersion) -> bool:
        return self._rules_checker.versions[ip_version].running

    @expect(MessageType.BlockUnblockList)
    async def process_block_unblock_list(self, message):
        """Block/unblock in batch in ipset"""
        # ipset are reloaded back from the db when rules recreated
        async with self._rules_checker.lock_rules_create_destroy:
            # remove ips with inactive ip versions
            lists = {
                "blocklist":
                    {
                        (ip, listname): properties
                        for (ip, listname), properties
                        in message["blocklist"].items()
                        if self._rules_running(IP.type_of(ip))},
                "unblocklist":
                    [
                        (ip, listname) for ip, listname
                        in message["unblocklist"]
                        if self._rules_running(IP.type_of(ip))
                    ]
                }
            num_tot = sum(map(len, lists.values()))

            async def block_unblock_many(IPSet, listname):
                """
                Add *listname*'s *blocklist* ips to *IPSet*.
                Remove *unblocklist* ips from *IPSet*.
                Remove non-*listname*'s *blocklist* ips from *IPSet*
                """
                ipset = IPSet()
                if not ipset.is_enabled():
                    logger.warning(
                        "%s ipset is disabled. Skip block/unblock.",
                        ipset.__class__.__qualname__
                    )
                    return
                to_block = (
                    (ip, properties)
                    for (ip, block_listname), properties
                    in lists['blocklist'].items()
                    if block_listname == listname
                )

                to_unblock = (
                    ip for ip, unblock_listname
                    in lists["unblocklist"]
                    if unblock_listname == listname)
                await ipset.restore(to_block, to_unblock)
            with timeit("Restoring to ipset "
                        "on synclist response %d ips" % num_tot,
                        logger):
                await block_unblock_many(IPSetIgnore, IPList.IGNORE)
                await block_unblock_many(IPSetGray, IPList.GRAY)
                await block_unblock_many(
                    IPSetGraySplashScreen, IPList.GRAY_SPLASHSCREEN
                )
                await block_unblock_many(IPSetBlack, IPList.BLACK)
                await block_unblock_many(IPSetWhite, IPList.WHITE)

    @expect(MessageType.FilesUpdated, files_type=WHITELISTS)
    async def process_global_whitelist_update(self, _):
        # whitelist update will be applied to ipset when rules/ipsets
        # are reloaded back from the db
        async with self._rules_checker.lock_rules_create_destroy:
            logger.info("Applying global white list update")
            for ipset in [IPSetStatic(), IPSetStaticRemoteProxy()]:
                if ipset.is_enabled():
                    await ipset.reset()

    @expect(MessageType.FilesUpdated, files_type=GEO)
    async def update_rules_and_ipset(self, _) -> None:
        async with self._rules_checker.lock_rules_create_destroy:
            logger.info("Updating ipset rules on geo ip update")
            for ip_version in ip_versions.enabled():
                await IPSetCountry().restore(ip_version)

    @expect(MessageType.UpdateCustomLists)
    async def process_custom_lists_update(self, _):

        async with self._rules_checker.lock_rules_create_destroy:
            await IPSetCustomBlacklist().reset()
            await IPSetCustomWhitelist().reset()
            logger.info("IPSet custom black and white lists are updated")

    @expect(MessageType.ConfigUpdate)
    @utils.log_error_and_ignore()
    async def process_config_update(self, message):
        """
        Recreate firewall rules and/or ipsets if corresponding settings change.
        """
        if isinstance(message["conf"], UserConfig):
            return  # do nothing for non-root config updates
        async with self._process_config_update_lock:
            return await self._on_config_update_unlocked(message)

    async def _on_config_update_unlocked(self, message):
        recreate = False
        current_status = (
            Webshield.ENABLE,
            Webshield.SPLASH_SCREEN,
            Webshield.PANEL_PROTECTION,
        )
        recreate_any_way = False
        if current_status != self._webshield_status:
            logger.info(
                "Webshield status (Webshield.ENABLE, Webshield.SPLASH_SCREEN) "
                "changed from %s to %s",
                self._webshield_status,
                current_status,
            )
            if Webshield.PANEL_PROTECTION != self._webshield_status[-1]:
                recreate_any_way = True
            self._webshield_status = current_status
            recreate = True
        current_port_blocking_mode = Firewall.port_blocking_mode
        if self._port_blocking_mode != current_port_blocking_mode:
            logger.info(
                "Ports blocking mode changed from %s to %s",
                self._port_blocking_mode,
                current_port_blocking_mode,
            )
            self._port_blocking_mode = current_port_blocking_mode
            recreate = True
        if recreate:  # recreate everything
            await self._rules_checker.recreate_rules_if_needed(
                recreate_any_way)
            logger.info('Firewall rules recreated due to ConfigUpdate')
        else:  # update just port blocking deny mode
            if await self._update_port_blocking_deny_mode_ipsets_if_needed():
                logger.info('Blocked ports deny mode updated on ConfigUpdate')

    async def _update_port_blocking_deny_mode_ipsets_if_needed(self):
        updated = False
        new_ports_blocking_values = self._get_port_blocking_deny_mode_values()
        if new_ports_blocking_values != self._port_blocking_deny_mode_values:
            logger.info(
                "Port blocking deny mode changed from %s to %s",
                _format_ports(self._port_blocking_deny_mode_values),
                _format_ports(new_ports_blocking_values),
            )
            self._port_blocking_deny_mode_values = new_ports_blocking_values
            async with self._rules_checker.lock_rules_create_destroy:
                sets = [
                    InputPortBlockingDenyModeIPSet(),
                    OutputPortBlockingDenyModeIPSet(),
                    IPSetNoRedirectPort()
                ]
                for ip_set in sets:
                    # FIREWALL.port_blocking_mode supports only ipv4
                    await ip_set.restore(IP.V4)
            updated = True
        return updated

    @expect(MessageType.StrategyChange)
    async def on_strategy_change(self, message: MessageType.StrategyChange):
        await self._rules_checker.recreate_rules_if_needed()
        logger.info('Firewall rules recreated due to StrategyChange %s',
                    message.strategy)


def _log_ipsets_mismatch(missing_ipsets, redundant_ipsets, log):
    """Report missing/redundant ipsets."""
    assert missing_ipsets or redundant_ipsets
    log(
        "Detected %s%s%s ipsets while ensuring ipsets/rules%s%s",
        "missing" * bool(missing_ipsets),
        "/" * bool(missing_ipsets and redundant_ipsets),
        "redundant" * bool(redundant_ipsets),
        f"; missing ipsets: {missing_ipsets}" * bool(missing_ipsets),
        f"; redundant ipsets: {redundant_ipsets}" * bool(redundant_ipsets),
    )


def _format_ports(ports):
    """Format ports for logging."""
    # note: the order is asserted in unit tests
    return json.dumps(
        dict(
            zip(
                "TCP_IN_IPV4 TCP_OUT_IPV4 UDP_IN_IPV4 UDP_OUT_IPV4".split(),
                ports,
            )
        )
    )
Cancel / Back
[ MAINHACK ]

Viewing / Editing File: lazy_init.py

Server Info

System Features

User Info
