^^&&^^

Edit File: webshield.py

"""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,
    Mode as WebshieldMode,
)
from defence360agent.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

"""
    current_mode = WebshieldMode.get()
    if WebshieldMode.wants_redirect(current_mode):
        redirect_map = webshield.port_redirect_map()
        dest_ports = webshield.redirected_to_webshield_ports(
            current_mode
        ) & 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 WebshieldMode.wants_redirect(current_mode):
        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"""
    current_mode = WebshieldMode.get()
    if WebshieldMode.wants_redirect(current_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)

# WARN: Module-based Apache/nginx must not be redirected to Webshield
    # and therefore their ports (80 & 443 by default) must have ACCEPT rules
    # for i360.ipv4.graylist and i360.ipv4.graysplashlist.
    # Otherwise, when IP is into one of these ipsets, Apache/nginx is unable
    # to response with SplashJS page - it will be dropped by next rule in chain
    # imunify360_log_gl.
    if current_mode == WebshieldMode.STANDALONE:
        module_based_ports = set()
    else:
        module_based_ports = get_module_based_ports()

yield FirewallRule(
        rule=FirewallRules.open_dst_ports_for_src_list(
            ipset_name,
            set(redirect_map[p] for p in dest_ports) | module_based_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
        ),
    )