­­­­­­­­­­­­­­­­­­
<!DOCTYPE html>
<html>
import asyncio
import datetime as DT
import logging
import time
from ipaddress import ip_network

import peewee

from defence360agent import utils
from defence360agent.contracts import messages, plugins
from defence360agent.model import simplification
from im360.api import ips
from im360.contracts import config
from im360.internals.core.firewall import is_nftables_backend
from im360.model import firewall
from im360.simple_rpc.lists import raise_acquired_validation_error
from im360.simple_rpc.resident_socket import send_to_socket

logger = logging.getLogger(__name__)
IP_LIST_UPDATE_TIMEOUT = 60  # seconds


class WhitelistCurrentUser(plugins.MessageSink):
    async def create_sink(self, loop):
        self._loop = loop

    @plugins.expect(messages.MessageType.CommandInvoke)
    async def whitelist_current_user_ip(self, message):
        try:
            ip = ip_network(message["params"].get("remote_addr"))
        except ValueError:
            return  # do nothing for invalid IP

        if ip.version == 6:
            ip = ip_network((int(ip.network_address), 64), strict=False)
        is_root = not message["params"].get("user")
        await add_to_whitelist(ip, full_access=is_root, loop=self._loop)


@raise_acquired_validation_error
async def _send_msg_to_socket(msg):
    return await send_to_socket(msg=msg, timeout=IP_LIST_UPDATE_TIMEOUT)


@utils.timed_cache(DT.timedelta(minutes=15))
async def add_to_whitelist(ip, full_access, *, loop=None):
    if loop is None:
        loop = asyncio.get_event_loop()
    logger.info("Adding %s to whitelist with ttl", ip)
    try:
        ip_obj = await simplification.run_in_executor(
            loop, lambda: firewall.IPList.get(ip=ip)
        )
    except peewee.DoesNotExist:
        pass
    else:
        if (
            ip_obj.listname == firewall.IPList.WHITE
        ) and not ip_obj.auto_whitelisted:
            # do not replace manually whitelisted
            return
        else:
            # Either expired BLACK/GRAY/GRAY_SPLASHSCREEN lists record
            # or auto-whitelisted, replacing
            if is_nftables_backend():
                # A group-scoped row is left alone: deleting it locally only
                # diverges from the group until the next group sync restores
                # it, and the whitelist entry added below takes precedence.
                if ip_obj.scope != firewall.IPList.SCOPE_GROUP:
                    # The resident deletes the row only on success, so a
                    # failed send leaves it intact and a retry retakes this
                    # branch. A local row makes it emit no group sync.
                    await _send_msg_to_socket(
                        {
                            "method": "IP_LISTS_UPDATE",
                            "action": "delete",
                            "purpose": (
                                firewall.IPListPurpose.listname2purpose(
                                    ip_obj.listname
                                )
                            ),
                            "items": [str(ip_obj.ip_network)],
                        }
                    )
            else:
                await ips.IPApi.unblock((ip_obj.ip_network,), ip_obj.listname)

    new_expiration = (
        0
        if config.AutoWhiteList.ttl_value() == 0
        else int(time.time() + config.AutoWhiteList.ttl_value())
    )

    comment = "IP auto-whitelisted with "
    if new_expiration == 0:
        comment += "no expiration"
    else:
        comment += "expiration date: {:%Y-%m-%d %H:%M:%S}".format(
            DT.datetime.fromtimestamp(new_expiration)
        )
    await send_to_socket(
        msg={
            "method": "IP_LISTS_UPDATE",
            "action": "add",
            "purpose": firewall.Purpose.WHITE,
            "items": {
                "items": [str(ip)],
                "expiration": new_expiration,
                "full_access": full_access,
                "comment": comment,
                "auto_whitelisted": True,
            },
        },
        timeout=IP_LIST_UPDATE_TIMEOUT,
    )
    logger.info("Added %s to whitelist with expiration %s", ip, new_expiration)