^^&&^^

Edit File: iptables.py

"""
Module for manipulating iptables rules
"""
import re
from itertools import groupby
from operator import itemgetter
from typing import List

from defence360agent.utils import CheckRunError, OsReleaseInfo, run
from defence360agent.utils.common import LooseVersion
from defence360agent.utils.validate import IP

from .base import (
    AbstractFirewall,
    FirewallBatchCommandError,
    FirewallCommandNotFoundError,
    FirewallError,
    FirewallTemporaryError,
)

_IPTABLES_LOCK_MSG = b"holding the xtables lock"
_IPTABLES_RESOURCE_MSG = b"Resource temporarily unavailable"
_BAD_LINE_RE = re.compile(
    rb"line (\d+) failed|Error occurred at line: (\d+)|"
    rb"[:] chain .*? is incompatible, use \'nft\' tool.\n\n"
)

class Iptables(AbstractFirewall):
    """
    Class wrapper for iptables executable
    """

EXIT_SUCCESS, EXIT_OTHER_PROBLEM, EXIT_PARAMETER_PROBLEM = 0, 1, 2

WAIT = "60"  # Waiting period is set to 60 seconds

def __init__(self, version, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.version = version
        self.restore_command = {
            IP.V4: ["iptables-restore"],
            IP.V6: ["ip6tables-restore"],
        }

# iptables-restore versions >= 1.4.21-18.0.1 have internal lock
        # to prevent concurrent access to iptables.
        # Wait option allows to wait for lock release
        # for a specific number of seconds avoiding crash
        # We can only check version here, but not build (18.0.1)
        # Build number is specified in spec-file to fullfill this requirement
        if not (OsReleaseInfo.DEBIAN & OsReleaseInfo.id_like()):
            # fix it later^: ubuntu iptables-1.6.0/iptables-restore
            #                has no support for '--wait'

if LooseVersion(self.version) >= LooseVersion("1.4.21"):
                for ipv in self.restore_command:
                    self.restore_command[ipv].extend(["--wait", self.WAIT])
        elif LooseVersion(self.version) >= LooseVersion("1.8.0"):
            for ipv in self.restore_command:
                self.restore_command[ipv].extend(["--wait", self.WAIT])

def has_rule(self, rule, table="filter", chain="INPUT", **kwargs):
        """
        Checks whether iptables rule exists

:param rule: rule parameters
        :param table: filter/nat/mangle
        :param chain: INPUT/PREROUTING/...
        :return: true, if rule exists
        """
        return dict(command="--check", table=table, chain=chain, rule=rule)

def has_chain(self, chain, table="filter"):
        """
        Checks whether chain exists

:param table: filter/nat/mangle
        :param chain: INPUT/PREROUTING/...
        :return: true, if rule exists
        """
        return dict(command="--list", table=table, chain=chain)

def append_rule(
        self, rule, table="filter", chain="INPUT", *args, **kwargs
    ):
        """
        Appends the rule.

:param rule: rule parameters
        :param table: filter/nat/mangle
        :param chain: INPUT/PREROUTING/...
        :return:
        """

return dict(command="--append", table=table, chain=chain, rule=rule)

def insert_rule(
        self, rule, table="filter", chain="INPUT", position=1, *args, **kwargs
    ):
        """
        Inserts the rule at specified position.

:param rule: rule parameters
        :param table: filter/nat/mangle
        :param chain: INPUT/PREROUTING/...
        :param position: rule number
        :return:
        """
        return dict(
            command="--insert",
            table=table,
            chain=chain,
            position=position,
            rule=rule,
        )

def delete_rule(
        self, rule, table="filter", chain="INPUT", *args, **kwargs
    ):
        """
        Deletes the rule

:param rule: rule parameters
        :param table: filter/nat/mangle
        :param chain: INPUT/PREROUTING/...
        :return:
        """
        return dict(command="--delete", table=table, chain=chain, rule=rule)

def create_chain(self, chain, table="filter"):
        """
        Creates new chain.

:param chain: name of chain
        :param table: table to add chain to
        :return:
        """
        return dict(command="--new", table=table, chain=chain)

def delete_chain(self, chain, table="filter"):
        """
        Deletes the chain.

:param chain: name of chain
        :param table: table to add chain to
        :return:
        """
        return dict(command="--delete-chain", table=table, chain=chain)

def flush_chain(self, chain, table="filter"):
        """
        Flushes rules in the chain.

:param chain: name of chain
        :param table: table to add chain to
        :return:
        """
        return dict(command="--flush", table=table, chain=chain)

def _prepare_output(self, records: List[dict]) -> str:
        output = []
        for table, actions in groupby(
            sorted(records, key=itemgetter("table")), key=itemgetter("table")
        ):
            output.append("*%s" % table)
            for action in actions:
                template = "{command} {chain} {position} {rule}"
                rule = action.get("rule", "")
                output.append(
                    template.format(
                        command=action["command"],
                        chain=action["chain"],
                        position=action.get("position", ""),
                        rule=" ".join(rule) if rule else "",
                    )
                )
            output.append("COMMIT")
        result = "\n".join(output) + "\n"
        return result

def _raise_for_output(
        self, code: int, out: bytes, stderr: bytes, commands: str
    ):
        cre = CheckRunError(
            returncode=code,
            cmd=self.restore_command[self.ip_version],
            output=out,
            stderr=stderr,
        )
        if _IPTABLES_LOCK_MSG in stderr:
            raise FirewallTemporaryError(
                self.ip_version, "Failed to acquire xtables lock"
            ) from cre
        if _IPTABLES_RESOURCE_MSG in stderr:
            raise FirewallTemporaryError(
                self.ip_version, _IPTABLES_RESOURCE_MSG.decode()
            ) from cre
        m = _BAD_LINE_RE.search(stderr)
        if m:
            n1, n2 = m.groups()
            n = int(n1 if n1 is not None else n2 if n2 is not None else 0)
            lines = commands.strip().split("\n")
            if n <= len(lines):
                err_msg = lines[n - 1]
            else:
                err_msg = "Failed to apply following rules:\n{}".format(
                    commands
                )
            raise FirewallBatchCommandError(self.ip_version, err_msg) from cre
        raise FirewallError(self.ip_version) from cre

async def _restore(self, output):
        try:
            code, out, err = await run(
                (
                    *self.restore_command[self.ip_version],
                    "--noflush",
                    "--verbose",
                ),
                input=output.encode(),
            )
        except FileNotFoundError as exc:
            raise FirewallCommandNotFoundError(self.ip_version, str(exc))
        if code == 0:
            return
        self._raise_for_output(code, out, err, output)

async def commit(self, records: List[dict]):
        await self._restore(self._prepare_output(records))