22

✏️ Editing: messages.py

from defence360agent import utils
from defence360agent.contracts.messages import (
    Accumulatable,
    Message,
    MessageList,
    Received,
    Reportable,
    ShortenReprListMixin,
)
from im360.contracts.message_pb2 import WebShieldData
from defence360agent.utils.validate import IP

HIGHEST_PRIORITY = 0

class StrategyChange(Message):
    """The message is generated when an IDS change is detected."""

DEFAULT_METHOD = "STRATEGY_CHANGE"
    PRIORITY = HIGHEST_PRIORITY

class SensorIncident(Message):
    """Single incident, e.g. user auth failed once"""

DEFAULT_METHOD = "INCIDENT"

class SensorIncidentList(MessageList, Reportable):
    """Aggregated incident list"""

DEFAULT_METHOD = "INCIDENT_LIST"

class UnreportableLocalIncidentList(MessageList):
    """Aggregate local incident list that are not reported to server"""

DEFAULT_METHOD = "LOCALINCIDENT_LIST"

class LocalIncidentList(MessageList, Reportable):
    """Aggregate local incident list - where no ip provided"""

DEFAULT_METHOD = "INCIDENT_LIST"

class SensorAlert(Message, Reportable):
    """Alert incident, e.g. user auth failures reached threshold"""

DEFAULT_METHOD = "ALERT"
    PRIORITY = 1

@classmethod
    def from_incident(cls, message):
        """When generate ALERT from INCIDENT change method"""
        new_message = message.copy()
        new_message["method"] = cls.DEFAULT_METHOD
        return cls(**new_message)

class ClientUnblock(Message, Reportable):
    DEFAULT_METHOD = "UNBLOCK"
    PRIORITY = HIGHEST_PRIORITY

class CaptchaEventList(ShortenReprListMixin, Message, Reportable):
    DEFAULT_METHOD = "CAPTCHA_LIST"

class CaptchaEvent(Accumulatable):
    DEFAULT_METHOD = "CAPTCHA"

FAILED = "FAILED"
    PASSED = "PASSED"
    REQUESTED = "REQUESTED"

LIST_CLASS = CaptchaEventList

@classmethod
    def from_parcel(cls, parcel):
        if parcel.websh.captcha == WebShieldData.NA:
            return None
        return cls(
            timestamp=parcel.timestamp,
            attackers_ip=parcel.ip,
            event=WebShieldData.Captcha.Name(parcel.websh.captcha),
            user_id=parcel.websh.user_id,
            plugin_id="captcha",
        )

class CaptchaDosAlert(Message):
    DEFAULT_METHOD = "CAPTCHA_DOS_ALERT"
    PRIORITY = 1

class SynclistResponse(Message, Received):
    DEFAULT_METHOD = "SYNCLIST"
    PRIORITY = 2

@utils.sync.timefun(
        action="SynclistResponse{str->IP.adopt_to_ipvX_network}"
    )
    def __init__(self, *args, **kwargs):
        """
        Do str -> Union[IPv4Network, IPv6Network]
        conversion only once per SynclistResponse message processing

:raise ValueError: if str keys to IPv4Network,
         IPv6Network conversion fails
        """
        super().__init__(*args, **kwargs)

for field in ["blocklist", "unblocklist"]:
            self[field] = {
                IP.adopt_to_ipvX_network(ip_str): dict_
                for ip_str, dict_ in self[field].items()
            }

@staticmethod
    def filter_blocklist(ips: dict, *, action_type: str):
        """Given [un]blocklist *ips* with their properties return ips
        matching given *action_type*.

ip without an action type set matches any action type.

"""
        return (
            ip
            for ip, p in (ips.items() if ips else [])
            if not p  # no properties
            or "action_type" not in p  # no action_type
            or p["action_type"] == action_type  # with given action_type
        )

class SynclistRequest(Message, Reportable):
    DEFAULT_METHOD = "SYNCLIST"

class BlockUnblockList(Message):
    """Used internally for block/unblock ip from lists
        {
            "blocklist": {(IPNetwork, "listname"): {"expiration": int}},
            "unblocklist": [(IPNetwork,"listname")] ,
        }

If ip is present in both lists: first unblock then block it
    (upsert semantics if applicable).

"""

DEFAULT_METHOD = "BLOCK_UNBLOCK"
    PRIORITY = HIGHEST_PRIORITY

class ProactiveQueueList(MessageList, Reportable):
    DEFAULT_METHOD = "PROACTIVE_QUEUE_LIST"

class RuleDisabled(Message, Reportable):
    """
    Rule disabled by customer
    """

DEFAULT_METHOD = "RULE_DISABLED"

class RuleEnabled(Message, Reportable):
    """
    Rule enabled back
    """

DEFAULT_METHOD = "RULE_ENABLED"

class ConfigSet(Message, Received):
    """Updates to the agent's config."""

RECEIVED_ACTIONS = ["CONFIG_SET"]

class UpdateCustomLists(Message):
    """
    Send message for class RealProtector for updating custom ip white and
    black list
    """

DEFAULT_METHOD = "UPDATE_CUSTOM_LISTS"

class IPListsUpdate(Message):
    """
    Send message for iplists plugin for updating IPs purposed lists
    """

DEFAULT_METHOD = "IP_LISTS_UPDATE"

class GroupIPSync(Message, Reportable):
    DEFAULT_METHOD = "GROUP_SYNC"

class GroupIPSyncPush(Message, Received):
    DEFAULT_METHOD = "GROUP_SYNC"

class EnduserPasswordReset(Message, Received):
    DEFAULT_METHOD = "ENDUSER_PASSWORD_RESET"

class WhitelistCacheUpdate(Message):
    """Receiveing this message indicates that resident part should clear whitelist cache"""

DEFAULT_METHOD = "WHITELIST_CACHE_UPDATE"

class IpsetUpdate(Message):
    """Receiveing this message indicates that resident part should check ipsets consistency"""

DEFAULT_METHOD = "IPSET_UPDATE"

class BlockedPortUpdate(Message):
    """Receiveing this message indicates that resident part should update blockedport"""

DEFAULT_METHOD = "BLOCKED_PORT_UPDATE"

class BlockedPortIPUpdate(Message):
    """Receiveing this message indicates that resident part should update blockedportip"""

DEFAULT_METHOD = "BLOCKED_PORT_IP_UPDATE"

MR SKB