MAINHACK

Viewing / Editing File: whitelist_current_user.py

File is not writable. Editing disabled.
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.model import firewall

logger = logging.getLogger(__name__)


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)


@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
            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 ips.IPApi.block(
        (ip,),
        firewall.IPList.WHITE,
        expiration=new_expiration,
        full_access=full_access,
        comment=comment,
        auto_whitelisted=True,
    )
    logger.info("Added %s to whitelist with expiration %s", ip, new_expiration)
Cancel / Back
[ MAINHACK ]

Viewing / Editing File: whitelist_current_user.py

Server Info

System Features

User Info
