Viewing File: /opt/imunify360/venv/versions/imunify-core-8.12.1-2/defence360agent/wordpress/incident_collector.py

"""Collector for WordPress CVE protection incidents."""

import logging
import os
import pwd
import stat as stat_module
import time
import re
from pathlib import Path
from collections import defaultdict

from defence360agent.model.wordpress import WPSite
from defence360agent.wordpress.cli import get_data_dir
from defence360agent.wordpress.incident_parser import IncidentFileParser
from defence360agent.model.wordpress_incident import (
    aggregate_incident_dicts,
    bulk_create_wordpress_incidents,
    build_incident_dict,
    country_reader,
)

logger = logging.getLogger(__name__)

#: Marks a file taken aside for processing. The plugin keeps appending to a
#: freshly created file under the original name, so nothing written during the
#: batch is lost, and a file left behind is retried on the next cycle. The
#: .php extension stays last: a webserver that only hands *.php to the
#: interpreter would serve any other extension as readable text.
PROCESSING_SUFFIX = ".processing.php"

#: A batch this old has failed every cycle since it was set aside. Retiring it
#: stops the retry from repeating forever and unblocks the same-hour file it
#: would otherwise keep out of collection.
QUARANTINE_AFTER_SECONDS = 15 * 60

#: Terminal names for a batch the collector must not pick up again. Neither
#: is matched by the pattern, and both keep .php last for the same reason
#: PROCESSING_SUFFIX does. A stored batch says so: its incidents are safe.
FAILED_SUFFIX = ".failed.php"
STORED_SUFFIX = ".stored.php"


class IncidentRateLimiter:
    """
    Rate limiter to prevent DoS attacks via incident flooding.

    Implements per-rule-per-IP rate limiting as per spec:
    - Maximum 100 incidents for each rule from the same IP within 15 minutes

    Memory-optimized implementation with bounded entry count using LRU eviction.
    """

    def __init__(
        self,
        max_incidents_per_rule_per_ip: int = 100,
        time_window_seconds: int = 900,  # 15 minutes
        max_unique_entries: int = 10000,  # Limit total unique (rule_id, IP) combinations
    ):
        """
        Initialize the rate limiter.

        Args:
            max_incidents_per_rule_per_ip: Max incidents per rule per IP (default: 100)
            time_window_seconds: Time window in seconds (default: 900 = 15 minutes)
            max_unique_entries: Max unique (rule_id, IP) combinations to track (default: 10000)
        """
        self.max_per_rule_per_ip = max_incidents_per_rule_per_ip
        self.time_window = time_window_seconds
        self.max_unique_entries = max_unique_entries

        # Track incident timestamps: {(rule_id, ip): [timestamp1, timestamp2, ...]}
        self.incident_times = defaultdict(list)

        self.cleanup_interval = 60  # Clean up old records every minute
        self.last_cleanup = time.time()

    def _cleanup_old_records(self):
        """Remove records older than the time window and enforce max entries limit."""
        now = time.time()
        cutoff = now - self.time_window

        # Clean expired timestamps from all entries
        keys_to_delete = []
        for key, timestamps in self.incident_times.items():
            # Filter out timestamps older than the window
            recent = [ts for ts in timestamps if ts > cutoff]

            if recent:
                self.incident_times[key] = recent
            else:
                keys_to_delete.append(key)

        for key in keys_to_delete:
            del self.incident_times[key]

        # Enforce max unique entries limit using LRU eviction
        if len(self.incident_times) > self.max_unique_entries:
            # Find oldest entries (those with oldest timestamp)
            entries_by_age = sorted(
                self.incident_times.items(),
                key=lambda x: x[1][0] if x[1] else 0,
            )

            # Remove oldest 10% of entries to avoid frequent evictions
            num_to_remove = max(
                1,
                len(self.incident_times) - int(self.max_unique_entries * 0.9),
            )
            for key, _ in entries_by_age[:num_to_remove]:
                del self.incident_times[key]

            logger.warning(
                "Rate limiter exceeded max entries (%d), removed %d oldest"
                " entries",
                self.max_unique_entries,
                num_to_remove,
            )

        self.last_cleanup = now

    def check_rate_limit(
        self, rule_id: str, attacker_ip: str, pending: int = 0
    ) -> tuple[bool, str]:
        """
        Check if adding an incident would exceed rate limits.

        Args:
            rule_id: Rule identifier
            attacker_ip: IP address of the attacker
            pending: Incidents already accepted in the current batch but not
                recorded yet, so one file cannot exceed the limit on its own

        Returns:
            Tuple of (allowed: bool, reason: str)
        """
        # Periodic cleanup
        if time.time() - self.last_cleanup > self.cleanup_interval:
            self._cleanup_old_records()

        now = time.time()
        cutoff = now - self.time_window
        key = (rule_id, attacker_ip)

        # Lazy cleanup: remove expired entries on access
        if key in self.incident_times:
            timestamps = self.incident_times[key]
            # Filter out old timestamps
            recent = [ts for ts in timestamps if ts > cutoff]

            if recent:
                self.incident_times[key] = recent
                recent_count = len(recent)
            else:
                # All timestamps expired, remove entry
                del self.incident_times[key]
                recent_count = 0
        else:
            recent_count = 0

        # Check if limit exceeded
        recent_count += pending
        if recent_count >= self.max_per_rule_per_ip:
            window_minutes = self.time_window // 60
            return (
                False,
                (
                    f"Rate limit exceeded for rule {rule_id} from IP"
                    f" {attacker_ip}:"
                    f" {recent_count}/{self.max_per_rule_per_ip} within"
                    f" {window_minutes} minutes"
                ),
            )

        return True, "OK"

    def record_incident(self, rule_id: str, attacker_ip: str):
        """
        Record that an incident was added.

        Args:
            rule_id: Rule identifier
            attacker_ip: IP address
        """
        now = time.time()
        key = (rule_id, attacker_ip)

        # Create list if it doesn't exist, or append to existing
        if key not in self.incident_times:
            self.incident_times[key] = [now]
        else:
            # Limit list size to prevent unbounded growth
            timestamps = self.incident_times[key]
            if len(timestamps) >= self.max_per_rule_per_ip:
                # Remove oldest timestamp when at limit
                timestamps.pop(0)
            timestamps.append(now)


class IncidentCollector:
    """
    Collect and persist WordPress incidents from plugin incident files.
    """

    def __init__(self, rate_limiter: IncidentRateLimiter | None = None):
        """
        Initialize the incident collector.

        Args:
            rate_limiter: Optional rate limiter (creates default if not provided)
        """
        self.rate_limiter = rate_limiter or IncidentRateLimiter()
        self.parser = IncidentFileParser()
        #: Batches this process read and failed to clear. The agent idles out
        #: after minutes of quiet, so age alone would retire one that was
        #: never retried.
        self._failed: set[str] = set()

    async def collect_incidents_for_site(
        self,
        site: WPSite,
        delete_after_processing: bool = True,
    ) -> list:
        """
        Collect incidents from a single WordPress site.

        Args:
            site: WordPress site to collect incidents from
            ruleset_version: Version of the ruleset being used
            delete_after_processing: Whether to delete incident files after processing

        Returns:
            List of collected Incident objects
        """
        collected_incidents = []

        try:
            data_dir = await get_data_dir(site)
            logger.debug("Data directory for site %s: %s", site, data_dir)
            if not data_dir.exists():
                logger.debug("Data directory does not exist for site %s", site)
                return []

            incident_files = self._get_incident_files(data_dir)
            logger.debug(
                "Incident files for site %s: %s", site, incident_files
            )
            if not incident_files:
                logger.debug("No incident files found for site %s", site)
                return []

            logger.debug(
                "Found %d incident file(s) for site %s",
                len(incident_files),
                site,
            )

            username = self._get_site_username(site)

            for incident_file in incident_files:
                file_incidents = await self._process_file(
                    incident_file,
                    site,
                    username,
                    delete_after_processing,
                )
                collected_incidents.extend(file_incidents)

        except Exception as e:
            logger.error(
                "Error collecting incidents for site %s: %s",
                site,
                e,
            )

        logger.info(
            "Collected %d incident(s) for site %s",
            len(collected_incidents),
            site,
        )

        return collected_incidents

    async def collect_incidents_for_sites(
        self,
        sites: list[WPSite],
        delete_after_processing: bool = True,
    ) -> list:
        """
        Collect incidents from multiple WordPress sites.

        Args:
            sites: List of WordPress sites
            delete_after_processing: Whether to delete incident files after processing

        Returns:
            List of collected Incident objects
        """
        all_collected_incidents = []

        for site in sites:
            site_incidents = await self.collect_incidents_for_site(
                site,
                delete_after_processing,
            )
            all_collected_incidents.extend(site_incidents)

        if all_collected_incidents:
            logger.info(
                "Collected %d WordPress incident(s) from %d site(s)",
                len(all_collected_incidents),
                len(sites),
            )

        return all_collected_incidents

    @classmethod
    def _get_incident_files(cls, data_dir: Path) -> list[Path]:
        """
        Get all incident files in the incidents directory.

        Args:
            data_dir: Path to the imunify-security data directory

        Returns:
            List of incident file paths
        """
        incidents_dir = data_dir / "incidents"
        logger.debug(
            "Incidents directory for site %s: %s", data_dir, incidents_dir
        )
        if not incidents_dir.exists() or not incidents_dir.is_dir():
            logger.debug(
                "Incidents directory does not exist for site %s", data_dir
            )
            return []

        # Use lstat (not Path.is_file) to identify regular files without following symlinks.
        incident_files = []
        for f in incidents_dir.iterdir():
            try:
                st = os.lstat(f)
            except OSError:
                continue
            if stat_module.S_ISREG(st.st_mode) and cls._is_incident_file(f):
                incident_files.append(f)

        logger.debug(
            "Incident files for site %s: %s", data_dir, incident_files
        )
        return incident_files

    # Pattern for incident files: yyyy-mm-dd-hh.php, optionally taken aside
    _FILE_PATTERN = re.compile(
        r"^\d{4}-\d{2}-\d{2}-\d{2}(?:\.processing)?\.php$"
    )

    @classmethod
    def _is_incident_file(cls, file_path: Path) -> bool:
        """
        Check if a file is an incident file based on naming pattern.

        Args:
            file_path: Path to the file to check

        Returns:
            True if file matches pattern yyyy-mm-dd-hh.php, with or without
            the suffix marking a batch left behind by an earlier cycle
        """
        return bool(cls._FILE_PATTERN.match(file_path.name))

    async def _process_file(
        self,
        incident_file,
        site: WPSite,
        username: str | None,
        delete_after_processing: bool,
    ) -> list:
        try:
            if delete_after_processing:
                incident_file = self._take_aside(incident_file)
                if incident_file is None:
                    return []

            incidents = self.parser.parse_file(incident_file)

            if incidents is None:
                self._failed.add(str(incident_file))
                return []

            if not incidents:
                logger.warning(
                    "No valid incidents in file %s",
                    incident_file.name,
                )
                if delete_after_processing:
                    self._discard(incident_file)

                return []

            logger.debug(
                "Parsed %d incident(s) from %s for site %s",
                len(incidents),
                incident_file.name,
                site,
            )

            collected_incidents = self._process_file_incidents(
                incidents,
                site,
                username,
                incident_file.name,
            )

            if delete_after_processing:
                self._discard(incident_file)

            return collected_incidents

        except Exception as e:
            if delete_after_processing:
                self._failed.add(str(incident_file))
            logger.error(
                "Error processing incident file %s for site %s: %s",
                incident_file.name,
                site,
                e,
            )
            return []

    def _take_aside(self, incident_file: Path) -> Path | None:
        """Move the file out of the plugin's way before reading it.

        Returns None when an earlier batch is still pending under the aside
        name; that batch is processed in its own turn and the fresh file waits
        for the next cycle rather than overwriting it.
        """
        if incident_file.name.endswith(PROCESSING_SUFFIX):
            if self._quarantine(incident_file):
                return None
            return incident_file

        aside = incident_file.with_name(
            incident_file.name[: -len(".php")] + PROCESSING_SUFFIX
        )
        if aside.exists() and self._pending(aside):
            if not self._quarantine(aside):
                return None

        incident_file.rename(aside)
        # rename keeps the plugin's mtime, so stamp the file to date the
        # attempt rather than the last write to it.
        os.utime(aside, None, follow_symlinks=False)
        return aside

    @staticmethod
    def _pending(aside: Path) -> bool:
        """Whether an aside still holds a batch waiting to be stored.

        Anything the site put there that is not a regular file is not one,
        and the rename replaces it.
        """
        try:
            st = os.lstat(aside)
        except OSError:
            return False
        return stat_module.S_ISREG(st.st_mode) and st.st_size > 0

    def _retire(self, path: Path, suffix: str) -> Path | None:
        """Give a batch a name the collector will not pick up again.

        Renaming touches the name, never what it points at, so it stays safe
        in a directory the site owns.
        """
        stem = path.name
        for known in (PROCESSING_SUFFIX, ".php"):
            if stem.endswith(known):
                stem = stem[: -len(known)]
                break

        retired = path.with_name(stem + suffix)
        try:
            path.rename(retired)
        except OSError as e:
            logger.error("Failed to retire %s: %s", path.name, e)
            return None

        self._failed.discard(str(path))
        return retired

    def _quarantine(self, aside: Path) -> bool:
        """Retire an aside this process has read and still failed to clear."""
        if str(aside) not in self._failed:
            return False

        try:
            age = time.time() - os.lstat(aside).st_mtime
        except OSError:
            return True

        if age < QUARANTINE_AFTER_SECONDS:
            return False

        retired = self._retire(aside, FAILED_SUFFIX)
        if retired is None:
            return False

        logger.error(
            "Gave up on %s after %d seconds, kept it as %s",
            aside.name,
            age,
            retired.name,
        )
        return True

    def _discard(self, incident_file: Path) -> bool:
        """Delete a stored batch, reporting whether it is gone.

        Emptying one we cannot delete would mean writing through a path the
        site owns, so it is retired under a name we never collect instead.
        """
        try:
            incident_file.unlink(missing_ok=True)
            self._failed.discard(str(incident_file))
            return True
        except OSError as e:
            logger.error("Failed to delete %s: %s", incident_file.name, e)

        # Its incidents are stored: reading the file again would add to their
        # counts, so retire it now rather than leaving it to be picked up.
        retired = self._retire(incident_file, STORED_SUFFIX)
        if retired is None:
            self._failed.add(str(incident_file))
        else:
            logger.warning(
                "Could not delete %s, kept the stored batch as %s",
                incident_file.name,
                retired.name,
            )
        return False

    def _get_site_username(self, site: WPSite) -> str | None:
        try:
            user_info = pwd.getpwuid(site.uid)
            return user_info.pw_name
        except Exception as e:
            logger.error(
                "Failed to get username for uid=%d, site %s: %s",
                site.uid,
                site,
                e,
            )
            return None

    def _process_file_incidents(
        self,
        incidents: list[dict],
        site: WPSite,
        username: str | None,
        incident_file_name: str,
    ) -> list:
        incidents_to_insert = []
        accepted: defaultdict = defaultdict(int)
        dropped_count = 0

        # Prepare all incidents for bulk insertion
        with country_reader() as geo_reader:
            for incident in incidents:
                rule_id = incident.get("rule_id", "unknown")
                attacker_ip = incident.get("REMOTE_ADDR") or incident.get(
                    "attacker_ip", "unknown"
                )

                allowed, reason = self.rate_limiter.check_rate_limit(
                    rule_id,
                    attacker_ip,
                    pending=accepted[(rule_id, attacker_ip)],
                )

                if not allowed:
                    logger.warning(
                        "Rate limit exceeded for site %s: %s",
                        site,
                        reason,
                    )
                    dropped_count += 1
                    continue

                # Prepare incident data for bulk insert
                site_info = {
                    "domain": site.domain,
                    "site_path": site.docroot,
                    "username": username,
                    "user_id": site.uid,
                }
                incident_data = build_incident_dict(
                    incident, site_info, geo_reader=geo_reader
                )

                incidents_to_insert.append(incident_data)
                accepted[(rule_id, attacker_ip)] += 1

        if not incidents_to_insert:
            logger.info(
                "Processed file %s: 0 stored, 0 aggregated, %d dropped",
                incident_file_name,
                dropped_count,
            )
            return []

        aggregated = aggregate_incident_dicts(incidents_to_insert)

        try:
            bulk_create_wordpress_incidents(aggregated)
        except Exception:
            logger.error(
                "Failed to store %d incident(s) from %s, keeping the file"
                " for the next cycle",
                len(incidents_to_insert),
                incident_file_name,
                exc_info=True,
            )
            raise

        # Only a stored incident spends rate-limit budget; a kept file is
        # retried next cycle and must not be throttled away unstored.
        for (rule_id, attacker_ip), count in accepted.items():
            for _ in range(count):
                self.rate_limiter.record_incident(rule_id, attacker_ip)

        logger.info(
            "Processed file %s: %d stored, %d aggregated, %d dropped",
            incident_file_name,
            len(aggregated),
            len(incidents_to_insert) - len(aggregated),
            dropped_count,
        )

        return aggregated
Back to Directory File Manager