"""
This module contains utilities to work with iptables to block SMTP
traffic on the server akin to how CSF does it.
"""

import grp
import itertools
import logging
import pwd
from contextlib import suppress
from typing import Generator, List, Tuple, NamedTuple

import im360.subsys.csf as csf
from defence360agent.utils import (
    Singleton,
    async_lru_cache,
)
from im360.contracts.config import Protector
from im360.contracts.config import SMTPBlocking as SMTPConfig
from im360.contracts.config import UnifiedAccessLogger
from im360.internals.core import ip_versions
from im360.internals.core.firewall import (
    FirewallRules,
    firewall_logging_enabled,
    get_firewall,
    is_nat_available,
    rule_logging_enabled,
)
from im360.internals.core.firewall.base import (
    FirewallError,
    FirewallBatchCommandError,
)
from im360.subsys.panels import hosting_panel


__all__ = [
    "sync_rules_for_all_versions",
    "reset_rules_for_all_versions",
    "is_SMTP_blocking_supported",
    "read_SMTP_settings",
    "get_active_settings_list",
    "conflicts_exist",
]

from defence360agent.utils.validate import IPVersion

TableState = NamedTuple(
    "TableState",
    [("chain_exists", bool), ("chain_referenced", bool), ("rules_ok", bool)],
)
SMTPSettings = NamedTuple(
    "SMTPSettings",
    [
        ("enabled", bool),
        ("ports", list),
        ("allow_users", set),
        ("allow_groups", set),
        ("allow_local", bool),
        ("redirect", bool),
    ],
)
CAPTURE_CSF_LOCK = True
CSF_LOCK_TIMEOUT = 300  # seconds

logger = logging.getLogger(__name__)

FILTER = "filter"
NAT = "nat"


async def _true_on_success(firewall, commands: List[dict]):
    """
    Use a non-zero return code of iptables as an indication of a failed check.

    An empty check command list is treated as a failure.

    Note: Should be called as close to the public functions as possible,
    since it has side effects.
    """

    if not commands:
        return False
    try:
        await firewall.commit(commands)
    except FirewallError:
        return False
    else:
        return True


def _get_uids(usernames):
    """Obtain UIDs of specified users skipping non-existing ones."""
    for user in usernames:
        try:
            yield pwd.getpwnam(user).pw_uid
        except KeyError:
            logger.warning("UNIX user %s does not exist", user)
            continue


def _get_gids(groups):
    """Obtain GIDs of specified groups skipping non-existing ones."""
    for group in groups:
        try:
            yield grp.getgrnam(group).gr_gid
        except KeyError:
            logger.warning("UNIX group %s does not exist", group)
            continue


class SMTPBlocking(metaclass=Singleton):
    """
    This class is used to synchronise iptables rules related to outgoing
    SMTP traffic blocking with SMTP_BLOCKING section of imunify config.
    """

    IM360_SMTP_CHAIN = "OUTPUT_imunify360_SMTP"
    IM360_SMTP_TARGET_RULE = ("-j", IM360_SMTP_CHAIN)

    def __init__(self, ip_version: IPVersion):
        self.ip_version = ip_version
        self._candidate_settings = None
        self.active_settings = None
        self.rules_were_reset = False
        self._hosting_panel = hosting_panel.HostingPanel()

    def _get_filter_smtp_rules(self) -> List[Tuple[str, ...]]:
        """
        Return a list of rules that should be used in
        OUTPUT_imunify360_SMTP chain.

        These can either be installed using append_rule / insert_rule or
        checked using has_rule methods of the firewall interface.
        """

        if not self._candidate_settings.ports:
            return []

        rules = []  # type: List[Tuple[str, ...]]
        common_args = (
            "-p",
            "tcp",
            "-m",
            "multiport",
            "--dports",
            ",".join(str(p) for p in self._candidate_settings.ports),
        )

        if self._candidate_settings.allow_local:
            rules.append(
                (*common_args, "-o", "lo", "-j", FirewallRules.ACCEPT)
            )

        # Allow root to send out anything.
        rules.append(
            (
                *common_args,
                "-m",
                "owner",
                "--uid-owner",
                "0",
                "-j",
                FirewallRules.ACCEPT,
            )
        )

        rules.extend(
            (
                *common_args,
                "-m",
                "owner",
                "--uid-owner",
                str(uid),
                "-j",
                FirewallRules.ACCEPT,
            )
            for uid in _get_uids(
                itertools.chain(
                    self._candidate_settings.allow_users,
                    self._hosting_panel.smtp_allow_users,
                )
            )
        )

        rules.extend(
            (
                *common_args,
                "-m",
                "owner",
                "--gid-owner",
                str(gid),
                "-j",
                FirewallRules.ACCEPT,
            )
            for gid in _get_gids(self._candidate_settings.allow_groups)
        )

        if firewall_logging_enabled() and rule_logging_enabled(
            UnifiedAccessLogger.SMTP
        ):
            rules.append(
                (
                    *common_args,
                    *FirewallRules.compose_rule(
                        action=FirewallRules.nflog_action(
                            group=FirewallRules.nflog_group(self.ip_version),
                            prefix=UnifiedAccessLogger.SMTP,
                        )
                    ),
                )
            )

        if not (
            self._candidate_settings.redirect
            and self._candidate_settings.allow_local
            and is_nat_available(self.ip_version)
        ):
            rules.append(
                (
                    *common_args,
                    "-j",
                    FirewallRules.REJECT,
                    "--reject-with",
                    "icmp{}-port-unreachable".format(
                        "6" if self.ip_version == ip_versions.IP.V6 else ""
                    ),
                )
            )

        return rules

    def _get_nat_smtp_rules(self) -> List[Tuple[str, ...]]:
        """
        Return a list of rules that should be used in
        OUTPUT_imunify360_SMTP chain in nat table.

        These can either be installed using append_rule / insert_rule or
        checked using has_rule methods of the firewall interface.
        """

        if not (
            self._candidate_settings.ports
            and self._candidate_settings.allow_local
            and self._candidate_settings.redirect
        ):
            return []

        rules = []  # type: List[Tuple[str, ...]]
        common_args = (
            "-p",
            "tcp",
            "-m",
            "multiport",
            "--dports",
            ",".join(str(p) for p in self._candidate_settings.ports),
        )

        rules.append((*common_args, "-o", "lo", "-j", FirewallRules.RETURN))

        # Allow root to send out anything.
        rules.append(
            (
                *common_args,
                "-m",
                "owner",
                "--uid-owner",
                "0",
                "-j",
                FirewallRules.RETURN,
            )
        )

        rules.extend(
            (
                *common_args,
                "-m",
                "owner",
                "--uid-owner",
                str(uid),
                "-j",
                FirewallRules.RETURN,
            )
            for uid in _get_uids(
                itertools.chain(
                    self._candidate_settings.allow_users,
                    self._hosting_panel.smtp_allow_users,
                )
            )
        )

        rules.extend(
            (
                *common_args,
                "-m",
                "owner",
                "--gid-owner",
                str(gid),
                "-j",
                FirewallRules.RETURN,
            )
            for gid in _get_gids(self._candidate_settings.allow_groups)
        )

        rules.append((*common_args, "-j", FirewallRules.REDIRECT))

        return rules

    def _get_smtp_rules_for(self, table: str) -> list:
        if table == FILTER:
            return self._get_filter_smtp_rules()
        if table == NAT:
            return self._get_nat_smtp_rules()
        return []

    def _im360_chain_exists(self, table, firewall):
        return [firewall.has_chain(table=table, chain=self.IM360_SMTP_CHAIN)]

    def _im360_chain_referenced(self, table, firewall):
        return [
            firewall.has_rule(
                table=table, chain="OUTPUT", rule=self.IM360_SMTP_TARGET_RULE
            )
        ]

    def _im360_rules_ok(self, table, firewall) -> List[dict]:
        """
        Check if SMTP rules in Imunify chain are in accord with new settings.
        """
        check_commands = [
            *(
                firewall.has_rule(
                    table=table, chain=self.IM360_SMTP_CHAIN, rule=rule
                )
                for rule in self._get_smtp_rules_for(table)
            )
        ]
        return check_commands

    def _reset_commands(
        self, table, firewall
    ) -> Generator[List[dict], None, None]:
        """
        Return commands that will ensure no OUTPUT blocking on Imunify part.

        Since the possible errors need to be suppressed we yield
        commands in batches, each of which can only contain one
        error-prone command.
        """
        yield [
            firewall.delete_rule(
                table=table, chain="OUTPUT", rule=self.IM360_SMTP_TARGET_RULE
            )
        ]
        yield [
            firewall.flush_chain(table=table, chain=self.IM360_SMTP_CHAIN),
            firewall.delete_chain(table=table, chain=self.IM360_SMTP_CHAIN),
        ]

    def _should_create_rules(self, table_state: TableState) -> bool:
        """Check whether rules need to be recreated."""

        active = self.active_settings
        new = self._candidate_settings

        if active is None:
            return True

        return (
            not table_state.chain_exists
            or not table_state.rules_ok
            or (active.allow_local and not new.allow_local)
            or (active.redirect != new.redirect)
            or bool(active.allow_users - new.allow_users)
            or bool(active.allow_groups - new.allow_groups)
        )

    def _sync_commands(self, table: str, table_state: TableState, firewall):
        """
        Return commands that will ensure firewall rules are in accord
        with settings.
        """
        commands = []
        if not table_state.chain_exists:
            commands.append(
                firewall.create_chain(self.IM360_SMTP_CHAIN, table=table)
            )
        if not table_state.chain_referenced:
            commands.append(
                firewall.insert_rule(
                    table=table,
                    chain="OUTPUT",
                    rule=self.IM360_SMTP_TARGET_RULE,
                )
            )

        if self._should_create_rules(table_state):
            commands.append(
                firewall.flush_chain(self.IM360_SMTP_CHAIN, table=table)
            )
            commands.extend(
                firewall.append_rule(
                    table=table, chain=self.IM360_SMTP_CHAIN, rule=rule
                )
                for rule in self._get_smtp_rules_for(table)
            )
        else:
            if not self._candidate_settings.ports:
                commands.append(
                    firewall.flush_chain(self.IM360_SMTP_CHAIN, table=table)
                )

        return commands

    async def _reset_rules_in_table(self, table: str, firewall):
        for batch in self._reset_commands(table, firewall):
            with suppress(FirewallBatchCommandError):
                await firewall.commit(batch)
        logger.info("SMTP Rules in table '%s' have been reset", table)

    async def _sync_rules_in_table(self, table: str, firewall):
        chain_exists = await _true_on_success(
            firewall, self._im360_chain_exists(table, firewall)
        )
        if not self._candidate_settings.enabled:
            if chain_exists:
                await self._reset_rules_in_table(table, firewall)
            return

        table_state = TableState(
            chain_exists=chain_exists,
            chain_referenced=await _true_on_success(
                firewall, self._im360_chain_referenced(table, firewall)
            ),
            rules_ok=await _true_on_success(
                firewall, self._im360_rules_ok(table, firewall)
            ),
        )

        commands = self._sync_commands(table, table_state, firewall)
        if commands:
            await firewall.commit(commands)
            logger.info(
                "SMTP settings have been synced with the rules in table '%s'",
                table,
            )

    async def reset_rules(self, firewall=None) -> None:
        """Ensure no OUTPUT blocking on Imunify part."""
        firewall = firewall or await get_firewall(self.ip_version)
        await self._reset_rules_in_table(FILTER, firewall)
        if is_nat_available(self.ip_version):
            await self._reset_rules_in_table(NAT, firewall)
        self.rules_were_reset = True

    async def sync_rules(self, new_settings) -> None:
        """Ensure iptables rules are in accord with settings."""
        firewall = await get_firewall(self.ip_version)
        self._candidate_settings = new_settings
        await self._sync_rules_in_table(FILTER, firewall)
        if is_nat_available(self.ip_version):
            await self._sync_rules_in_table(NAT, firewall)
        # Active settings should only be updated after we finished
        # modifying all the tables.
        self.active_settings = self._candidate_settings
        self.rules_were_reset = False


async def _reset_rules_for_ip_versions(ip_versions_to_reset: List):
    for version in ip_versions_to_reset:
        if not SMTPBlocking(version).rules_were_reset:
            await SMTPBlocking(version).reset_rules()


async def reset_rules_for_all_versions():
    """
    Mainly used for `SMTPBlocker` plugin shutdown.
    """
    ip_versions_to_reset = [
        version
        for version in ip_versions.all()
        if not SMTPBlocking(version).rules_were_reset
    ]
    if ip_versions_to_reset:
        await _reset_rules_for_ip_versions(ip_versions_to_reset)


async def sync_rules_for_all_versions(new_settings):
    """
    Used whenever there is a need to check compatibility between Imunify
    config and currently used SMTP blocking iptables rules.
    """
    async with Protector.RULE_EDIT_LOCK:
        for version in ip_versions.enabled():
            await SMTPBlocking(version).sync_rules(new_settings)


@async_lru_cache(maxsize=1)
async def is_SMTP_blocking_supported():
    """Check if iptables has xt_owner module."""
    async with Protector.RULE_EDIT_LOCK:
        # IPv4 and IPv6 share the same xt_owner module
        firewall = await get_firewall(ip_versions.IP.V4)
        try:
            await firewall.commit(
                [
                    firewall.insert_rule(
                        chain="OUTPUT", rule=FirewallRules.smtp_test_rule()
                    )
                ],
            )
        except FirewallError:
            return False

        try:
            await firewall.commit(
                [
                    firewall.delete_rule(
                        chain="OUTPUT", rule=FirewallRules.smtp_test_rule()
                    )
                ],
            )
        except FirewallError as err:
            logger.warning(
                "Something went wrong"
                " during the removal of the SMTP test rule: %s",
                err,
            )
    return True


def read_SMTP_settings() -> SMTPSettings:
    """Return current settings from Imunify config."""
    return SMTPSettings(
        enabled=SMTPConfig.ENABLED,
        ports=SMTPConfig.PORTS,
        allow_users=set(SMTPConfig.ALLOW_USERS),
        allow_groups=set(SMTPConfig.ALLOW_GROUPS),
        allow_local=SMTPConfig.ALLOW_LOCAL,
        redirect=SMTPConfig.REDIRECT,
    )


def get_active_settings_list() -> list:
    """
    Return the latest applied SMTP settings.

    Used to compare with the settings from config file.
    """
    active_settings_list = []
    for version in ip_versions.enabled():
        active_settings_list.append(SMTPBlocking(version).active_settings)
    return active_settings_list


async def conflicts_exist() -> bool:
    """
    Return True if any other SMTP blocking features is active
    """

    panel_SMTP_conflict = (
        hosting_panel.HostingPanel().get_SMTP_conflict_status()
    )
    return any((await csf.is_SMTP_block_enabled(), panel_SMTP_conflict))
