LittleDemon WebShell

"""Webshield related iptables rules."""
from typing import AbstractSet, Iterator, Mapping

from im360.internals.core import firewall
from im360.internals.core.firewall import FirewallRules
from im360.subsys import webshield
from im360.subsys.webshield_mode import (
    get_module_based_ports,
    is_standalone_mode,
)
from im360.utils.validate import IPVersion

from .port import redirect_port_rules
from .types_ import FirewallRule, WebshieldRuleBuilder


def rules(
    ipset_name: str, ip_version: IPVersion, gray_rules: WebshieldRuleBuilder
) -> Iterator[FirewallRule]:
    """Yield iptables *ip_version* rules for *ipset_name* using
    *gray_rules* builder.

    This is intended for ipsets that require webshield to implement
    their behavior e.g., to show captcha, splashscreen to the ips from
    the ipset

    """
    if is_standalone := is_standalone_mode():
        redirect_map = webshield.port_redirect_map()
        dest_ports = webshield.redirected_to_webshield_ports() & set(
            redirect_map
        )
    else:
        dest_ports = get_module_based_ports()
        redirect_map = {port: port for port in dest_ports}

    yield from _redirect_rules(
        ipset_name, ip_version, redirect_map, dest_ports, gray_rules
    )

    if not is_standalone:
        if firewall.is_nat_available(ip_version):
            yield from gray_rules.logdrop_chain_rules(ipset_name)
        else:
            yield from gray_rules.drop_tproxy_rules(ipset_name)
        return

    if firewall.is_nat_available(ip_version):
        yield from gray_rules.logdrop_chain_rules(ipset_name)
        yield from redirect_port_rules(
            ipset_name,
            dest_ports,
            redirect_map,
            FirewallRules.NAT,
            FirewallRules.redirect_to_captcha,
        )
    else:
        # What we are doing if we encounter with centos 6 and ipv6?

        # Full description is available into:
        # https://cloudlinux.atlassian.net/browse/DEF-2898
        # https://access.redhat.com/solutions/311493

        # Mark traffic to http, https hosts and from graylist ip
        yield from redirect_port_rules(
            ipset_name,
            dest_ports,
            redirect_map,
            FirewallRules.MANGLE,
            FirewallRules.redirect_to_captcha_via_tproxy,
        )
        yield from gray_rules.drop_tproxy_rules(ipset_name)


def _redirect_rules(
    ipset_name: str,
    ip_version: IPVersion,
    redirect_map: Mapping[int, int],
    dest_ports: AbstractSet[int],
    gray_rules: WebshieldRuleBuilder,
) -> Iterator[FirewallRule]:
    """Yield rules for to-be-redirected ports"""
    if is_standalone_mode():
        yield from check_access_to_webshield_ports_rules(
            ipset_name, set(redirect_map[p] for p in dest_ports)
        )
        yield from gray_rules.open_webshield_ports_for_localhost_rules(
            ip_version
        )
        yield from gray_rules.block_webshield_ports_rules(
            redirect_map, dest_ports
        )
        yield from gray_rules.redirect_panel_ports(ip_version)

    yield FirewallRule(
        rule=FirewallRules.open_dst_ports_for_src_list(
            ipset_name, set(redirect_map[p] for p in dest_ports)
        ),
    )


def check_access_to_webshield_ports_rules(
    ipset_name: str, dest_ports: AbstractSet[int]
) -> Iterator[FirewallRule]:
    yield FirewallRule(
        chain=FirewallRules.WEBSHIELD_PORTS_INPUT_CHAIN,
        rule=FirewallRules.open_dst_ports_for_src_list(
            ipset_name, dest_ports, policy=FirewallRules.RETURN
        ),
    )