^^&&^^

Edit File: libipset.py

"""Python API for calling ipset command-line utility."""

import asyncio
import os
import tempfile
from functools import lru_cache
from logging import getLogger
from typing import Iterable, Optional

from defence360agent.utils import (
    CheckRunError,
    await_for,
    check_run,
    readlines_from_cmd_output,
    retry_on,
)
from defence360agent.utils.common import DAY, rate_limit

#: One command (a string) that can be fed to `ipset restore` utility
IPSetRestoreCmd = str  # TypeAlias is 3.10+

logger = getLogger(__name__)
throttled_log_error = rate_limit(period=DAY, on_drop=logger.warning)(
    logger.error
)

HASH_IP = "hash:ip"
HASH_NET = "hash:net"
HASH_NET_PORT = "hash:net,port"
BITMAP_PORT = "bitmap:port"

# Datatypes
IPSET_NET, IPSET_PORT, IPSET_NET_PORT = "net", "port", "net,port"
IPSET_RESTORE_TIMEOUT = 15
#: http://ipset.netfilter.org/ipset.man.html
IPSET_TIMEOUT_MAX = 2147483  # seconds

_COMMANDS_LOG_SIZE_CUTOFF = 80 * 3  # keep ~3 full lines at most

@lru_cache(1)
def get_ipset_exe():
    """
    Returns the path to the ipset executable.
    :return: path to ipset executable
    """
    for path in [
        "/usr/sbin/ipset",
        "/usr/bin/ipset",
        "/sbin/ipset",
        "/bin/ipset",
    ] + [
        os.path.join(p, "ipset")
        for p in os.environ.get("PATH", "").split(os.pathsep)
    ]:
        # Check if the path exists and is a file and executable
        if os.path.isfile(path) and os.access(path, os.X_OK):
            return path
    # If no path is found return binary name
    logger.error(
        "Cannot find executable with the name 'ipset' in PATH %s",
        os.environ.get("PATH"),
    )
    return "ipset"

def _gen_ipset_cmd(call_args):
    return [get_ipset_exe()] + call_args

class IPSetError(RuntimeError):
    """Base class for libipset errors."""

class IgnoredIPSetKernelError(IPSetError):
    pass

class IPSetCannotBeDestroyedError(IPSetError):
    SIGNATURE = b"Set cannot be destroyed: it is in use by a kernel component"

class IPSetNotFoundError(IPSetError):
    SIGNATURE = b"The set with the given name does not exist"

class IPSetKernelPermittedError(IgnoredIPSetKernelError):
    SIGNATURE = b"Kernel error received: Operation not permitted"

class IPSetKernelSessionError(IgnoredIPSetKernelError):
    SIGNATURE = b"Cannot open session to kernel"

class IPSetKernelBusyError(IgnoredIPSetKernelError):
    SIGNATURE = b"Kernel error received: Device or resource busy"

class IPSetEXENotFoundError(IPSetError):
    SIGNATURE = b"Cannot find executable with the name 'ipset'"

class IPSetCmdBuilder:
    """Generate commands that can be passed to `ipset restore` utility."""

@staticmethod
    def get_add_cmd(ipset_name: str, entry: str) -> IPSetRestoreCmd:
        return f"add {ipset_name} {entry} -exist"

@staticmethod
    def get_delete_cmd(ipset_name: str, entry: str) -> IPSetRestoreCmd:
        return f"del {ipset_name} {entry} -exist"

@staticmethod
    def get_create_cmd(
        ipset_name: str,
        family,
        datatype=HASH_NET,
        timeout=0,
        maxelem=65536,
    ) -> IPSetRestoreCmd:
        return (
            f"create {ipset_name} {datatype} family {family} timeout"
            f" {timeout} maxelem {maxelem} -exist"
        )

@staticmethod
    def get_create_list_set_cmd(
        ipset_name: str, size: int = 8
    ) -> IPSetRestoreCmd:
        return f"create {ipset_name} list:set size {size} -exist"

@staticmethod
    def get_destroy_cmd(ipset_name: str) -> IPSetRestoreCmd:
        return f"destroy {ipset_name}"

@staticmethod
    def get_flush_cmd(ipset_name: str) -> IPSetRestoreCmd:
        return f"flush {ipset_name}"

def raise_ipset_error_if_matched(exc, msg):
    for exc_cls in [
        IPSetCannotBeDestroyedError,
        IPSetNotFoundError,
        IPSetKernelSessionError,
        IPSetKernelBusyError,
        IPSetKernelPermittedError,
    ]:
        if exc_cls.SIGNATURE in exc.stderr:
            raise exc_cls(msg) from exc
    raise exc

async def _run_ipset(name, command, *args, **kwargs):
    """
    :param name: ipset name, None if we can't get set name
    :param command: ipset will run as 'IPSET_EXEC <command>'
    """
    command = _gen_ipset_cmd(command)
    try:
        return await check_run(command, *args, **kwargs)
    except CheckRunError as e:
        raise_ipset_error_if_matched(
            e,
            f"Error '{e.stderr}' occurs when executing '{command}' "
            f"command for '{name}'",
        )

def prepare_ipset_command(cmd, name, item, timeout=0):
    if cmd == "add":
        if timeout > IPSET_TIMEOUT_MAX:
            throttled_log_error(
                "Wrong timeout: %s %s %s %s; clipped to %s",
                cmd,
                name,
                item,
                timeout,
                IPSET_TIMEOUT_MAX,
            )
            timeout = IPSET_TIMEOUT_MAX
        return [cmd, name, str(item), "timeout", str(timeout), "-exist"]
    elif cmd == "del":
        return [cmd, name, str(item), "-exist"]
    else:
        raise NotImplementedError(
            "Method with action {} not implemented".format(cmd)
        )

async def add_item(name, item, timeout):
    """
    Adds entry into existing set of ipset

:param str name: name of set from ipset
    :param str item: IP v4 address
    :param int timeout: relative timeout in seconds
    :return:
    """

command = prepare_ipset_command("add", name, item, timeout)

await _run_ipset(name, command)

async def delete_item(name, item):
    """
    Removes entry from existing set of ipset

:param str name: name of set from ipset
    :param str item: IP v4 address
    :return:
    """

command = prepare_ipset_command("del", name, item)

await _run_ipset(name, command)

# TODO: Refactor to avoid code duplication
async def create_hash_set(name, datatype=IPSET_NET, **options):
    """
    Creates hashset into ipset.
    :param name: name of the set
    :param datatype: type of stored data (ip, net, port, (net, port))
    :param options: options to command
    :return:
    """
    if not isinstance(name, str):
        raise TypeError(
            "{name} is {type_} but str expected".format(
                name=name, type_=type(name)
            )
        )
    datatypes = [IPSET_NET, IPSET_NET_PORT]
    if datatype not in datatypes:
        raise ValueError(
            "Datatype argument value should be in {datatypes}, "
            "but {datatype} received".format(
                datatypes=datatypes, datatype=datatype
            )
        )

set_type = "hash:" + datatype
    options = {k: str(v) for k, v in options.items()}

command = ["create", name, set_type]

command.extend(["family", options.get("family", "inet")])
    command.extend(["maxelem", options.get("maxelem", "65536")])
    command.extend(["timeout", options.get("timeout", "0")])
    command.append("-exist")

await _run_ipset(name, command)

async def create_bitmap_set(name, datatype=IPSET_PORT, **options):
    """
    Creates bitmapset into ipset.
    :param name: name of the set
    :param datatype: type of stored data (ip, net, port)
    :param options: options to command
    :return:
    """
    if not isinstance(name, str):
        raise TypeError(
            "{name} is {type_} but str expected".format(
                name=name, type_=type(name)
            )
        )
    datatypes = [IPSET_PORT]
    if datatype not in datatypes:
        raise ValueError(
            "Datatype argument value should be in {datatypes}, "
            "but {datatype} received".format(
                datatypes=datatypes, datatype=datatype
            )
        )
    set_type = "bitmap:" + datatype
    options = {k: str(v) for k, v in options.items()}

command = ["create", name, set_type]

command.extend(["range", options.get("range", "0-65535")])
    command.extend(["timeout", options.get("timeout", "0")])
    command.append("-exist")

await _run_ipset(name, command)

@retry_on(
    IPSetCannotBeDestroyedError, max_tries=3, on_error=await_for(seconds=3)
)
async def delete_set(name):
    """
    Removes set into ipset
    Removes rule into firewall-cmd wich links with new set of ipset

:param str name: name of set
    :return:
    """
    assert isinstance(name, str)

existing = await list_set()
    if name in existing:
        await _run_ipset(name, ["flush", name])
        await _run_ipset(name, ["destroy", name])

async def flush_set(name):
    """
    Removes ips from set

:param str name: name of set
    :return:
    """
    assert isinstance(name, str)

existing = await list_set()
    if name in existing:
        await _run_ipset(name, ["flush", name])

async def list_set():
    """
    Returns names of ipset sets

:return:
    """
    out = await _run_ipset(None, ["list", "-n", "-t"])
    out = out.decode().strip()  # type: str
    if out:
        return out.splitlines(keepends=False)
    else:
        return []

async def restore(
    lines: Iterable[IPSetRestoreCmd], name: Optional[str] = None
) -> bytes:
    """
    Run `ipset restore` command for bulk operations

:param lines: lines of input commands for `ipset restore` utility
    :param name: optional ipset name to be used in error reporting
    :return: the output of ipset command as bytes
    """
    with tempfile.TemporaryFile() as f:
        # note: can't use to_thread here because *lines* may invoke db ops
        #     DEF-15621 may fix it
        f.writelines(line.encode() + b"\n" for line in lines)
        f.flush()
        f.seek(0)
        try:
            async with asyncio.timeout(IPSET_RESTORE_TIMEOUT):
                return await _run_ipset(name, ["restore"], stdin=f)
        except Exception as e:
            # preserve specific IPSetError type
            Error = e.__class__ if isinstance(e, IPSetError) else IPSetError
            # add ipset commands to the error message
            file_size = os.fstat(f.fileno()).st_size
            f.seek(0)
            if file_size < _COMMANDS_LOG_SIZE_CUTOFF:
                commands = f.read()
            else:  # cut
                commands = f.read(_COMMANDS_LOG_SIZE_CUTOFF // 2)
                commands += b"..."
                f.seek(-_COMMANDS_LOG_SIZE_CUTOFF // 2, os.SEEK_END)
                commands += f.read()
            raise Error(
                "ipset restore failed. "
                f"Name: {name!r} "
                f"Reason: {e!r} "
                f"Commands: {commands!r}"
            ) from e

async def swap(set_name1, set_name2):
    await _run_ipset(None, ["swap", set_name1, set_name2])

async def get_ipset_count(setname: str) -> int:
    """Return the number of ips in the *setname* ipset."""
    command = _gen_ipset_cmd(["list", setname, "-terse"])
    try:
        async with asyncio.timeout(IPSET_RESTORE_TIMEOUT):
            async for line in readlines_from_cmd_output(command):
                if b"Number of entries:" in line:
                    return int(line.split(b":")[-1])
    except CheckRunError:
        pass

return await _get_ipset_count_wc(setname)

async def _get_ipset_count_wc(setname: str) -> int:
    command = f"ipset save {setname} | grep '^add' | wc -l"
    try:
        async with asyncio.timeout(IPSET_RESTORE_TIMEOUT):
            process = await asyncio.create_subprocess_shell(
                command,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
            )
            stdout, _ = await process.communicate()
            if process.returncode == 0:
                return int(stdout.decode().strip())
    except (ValueError, TimeoutError, CheckRunError):
        pass
    return 0

async def get_ipset_members(setname: str) -> set[str]:
    """Return the members of the *setname* ipset."""
    command = _gen_ipset_cmd(["list", setname])
    members = set()
    member_section = False
    try:
        async with asyncio.timeout(IPSET_RESTORE_TIMEOUT):
            async for line in readlines_from_cmd_output(command):
                if b"Members:" in line:
                    member_section = True
                    continue
                if member_section:
                    members.add(line.strip())
    except CheckRunError:
        pass
    return members