rengine/shared/definitions/notifications.py

from shared.definitions.vulnerabilities import (
    ALERT_SEVERITIES,
    SEVERITY_LABELS,
    SEVERITY_ORDER,
    Severity,
)
from shared.enums.notification import NotificationSeverity, NotificationType


def whois_enrichment_complete(
    success: int,
    failed: int,
    total: int,
) -> dict:
    if failed == 0:
        return {
            "type": NotificationType.TARGET,
            "severity": NotificationSeverity.SUCCESS,
            "title": "WHOIS Enrichment Complete",
            "message": (
                f"WHOIS lookup completed for {success} "
                f"{'target' if success == 1 else 'targets'}."
            ),
        }
    return {
        "type": NotificationType.TARGET,
        "severity": NotificationSeverity.WARNING,
        "title": "WHOIS Enrichment Complete",
        "message": (
            f"WHOIS lookup completed: {success} succeeded, "
            f"{failed} failed out of {total} targets."
        ),
    }


def whois_enrichment_failed(error: str) -> dict:
    return {
        "type": NotificationType.TARGET,
        "severity": NotificationSeverity.ERROR,
        "title": "WHOIS Enrichment Failed",
        "message": f"WHOIS enrichment task failed: {error}",
    }


def ripestat_enrichment_complete(
    success: int, failed: int, skipped: int, total: int
) -> dict:
    if failed == 0 and skipped == 0:
        severity = NotificationSeverity.SUCCESS
        title = "BGP Enrichment Complete"
        message = f"Successfully enriched {success} targets with BGP data."
    elif success > 0:
        severity = NotificationSeverity.WARNING
        title = "BGP Enrichment Partially Complete"
        message = (
            f"BGP enrichment finished: {success} succeeded, "
            f"{failed} failed, {skipped} skipped out of {total} targets."
        )
    else:
        severity = NotificationSeverity.ERROR
        title = "BGP Enrichment Failed"
        message = (
            f"BGP enrichment could not complete: "
            f"{failed} failed, {skipped} skipped out of {total} targets."
        )

    return {
        "type": NotificationType.TARGET,
        "severity": severity,
        "title": title,
        "message": message,
    }


def ripestat_enrichment_failed(error: str) -> dict:
    return {
        "type": NotificationType.TARGET,
        "severity": NotificationSeverity.ERROR,
        "title": "BGP Enrichment Failed",
        "message": f"BGP enrichment task failed: {error}",
    }


_SECONDS_PER_MINUTE = 60
_MINUTES_PER_HOUR = 60


def _scan_meta(scan_id: str) -> dict:
    return {"scan_id": str(scan_id), "url": f"/scans/{scan_id}"}


def _scan_duration(seconds: float | None) -> str:
    if not seconds or seconds < 0:
        return ""
    seconds = int(seconds)
    if seconds < _SECONDS_PER_MINUTE:
        return f" in {seconds}s"
    minutes, rem = divmod(seconds, _SECONDS_PER_MINUTE)
    if minutes < _MINUTES_PER_HOUR:
        return f" in {minutes}m {rem}s" if rem else f" in {minutes}m"
    hours, minutes = divmod(minutes, _MINUTES_PER_HOUR)
    return f" in {hours}h {minutes}m"


def scan_started(scan_id: str, target: str, engine: str) -> dict:
    return {
        "type": NotificationType.SCAN,
        "severity": NotificationSeverity.INFO,
        "title": "Scan Started",
        "message": f"Scan of {target} started with engine {engine}.",
        "metadata": _scan_meta(scan_id),
    }


_SCAN_COUNT_LABELS = {
    "subdomains_found": ("subdomain", "subdomains"),
    "ips_found": ("IP", "IPs"),
    "open_ports_found": ("open port", "open ports"),
    "http_assets_found": ("HTTP service", "HTTP services"),
    "vulnerabilities_found": ("vulnerability", "vulnerabilities"),
    "endpoints_found": ("endpoint", "endpoints"),
}


def scan_count_summary(counts: dict) -> str:
    parts = [
        f"{n} {singular if n == 1 else plural}"
        for col, (singular, plural) in _SCAN_COUNT_LABELS.items()
        if (n := counts.get(col, 0))
    ]
    return ", ".join(parts) if parts else "no assets"


def stage_count_summary(counts: dict) -> str:
    parts = [
        f"{n} {key.replace('_', ' ')}"
        for key, n in counts.items()
        if isinstance(n, int) and n and key != "excluded"
    ]
    return ", ".join(parts) if parts else "no results"


def scan_completed(
    scan_id: str,
    target: str,
    engine: str,
    counts: dict,
    duration_seconds: float | None,
) -> dict:
    return {
        "type": NotificationType.SCAN,
        "severity": NotificationSeverity.SUCCESS,
        "title": "Scan Completed",
        "message": (
            f"Scan of {target} ({engine}) completed"
            f"{_scan_duration(duration_seconds)} — {scan_count_summary(counts)}."
        ),
        "metadata": _scan_meta(scan_id),
    }


def scan_new_subdomains(scan_id: str, target: str, count: int) -> dict:
    return {
        "type": NotificationType.SCAN,
        "severity": NotificationSeverity.INFO,
        "title": "New Subdomains Discovered",
        "message": (
            f"{count} new {'subdomain' if count == 1 else 'subdomains'} "
            f"discovered on {target}."
        ),
        "metadata": _scan_meta(scan_id),
    }


def scan_new_services(scan_id: str, target: str, count: int, sensitive: int) -> dict:
    detail = (
        f" {sensitive} of them administrative or datastore "
        f"{'port' if sensitive == 1 else 'ports'}."
        if sensitive
        else ""
    )
    return {
        "type": NotificationType.SCAN,
        "severity": NotificationSeverity.WARNING
        if sensitive
        else NotificationSeverity.INFO,
        "title": "New Services Exposed",
        "message": (
            f"{count} {'service' if count == 1 else 'services'} on {target} "
            f"{'was' if count == 1 else 'were'} not open at the previous scan.{detail}"
        ),
        "metadata": _scan_meta(scan_id),
    }


def scan_failed(scan_id: str, target: str, engine: str, error: str) -> dict:
    return {
        "type": NotificationType.SCAN,
        "severity": NotificationSeverity.ERROR,
        "title": "Scan Failed",
        "message": f"Scan of {target} ({engine}) failed: {error[:300]}",
        "metadata": _scan_meta(scan_id),
    }


def scan_cancelled(scan_id: str, target: str, engine: str) -> dict:
    return {
        "type": NotificationType.SCAN,
        "severity": NotificationSeverity.WARNING,
        "title": "Scan Cancelled",
        "message": f"Scan of {target} ({engine}) was cancelled.",
        "metadata": _scan_meta(scan_id),
    }


def _severity_phrase(counts: dict) -> str:
    parts = [
        f"{counts[name]} {SEVERITY_LABELS[name].lower()}"
        for name in SEVERITY_ORDER
        if counts.get(name)
    ]
    return ", ".join(parts)


def scan_new_vulnerabilities(
    scan_id: str, target: str, counts: dict, kev: int, total: int
) -> dict:
    severe = sum(counts.get(name, 0) for name in ALERT_SEVERITIES)
    if counts.get(Severity.CRITICAL.value) or kev:
        severity = NotificationSeverity.ERROR
    elif severe:
        severity = NotificationSeverity.WARNING
    else:
        severity = NotificationSeverity.INFO
    detail = _severity_phrase(counts)
    kev_note = (
        f" {kev} {'is' if kev == 1 else 'are'} known to be exploited in the wild."
        if kev
        else ""
    )
    return {
        "type": NotificationType.VULNERABILITY,
        "severity": severity,
        "title": "New Vulnerabilities Found",
        "message": (
            f"{total} new {'finding' if total == 1 else 'findings'} on {target}"
            f"{f' — {detail}' if detail else ''}.{kev_note}"
        ),
        "metadata": {
            **_scan_meta(scan_id),
            "url": f"/scans/{scan_id}?tab=vulnerabilities",
        },
    }


def scan_vulnerability_coverage(scan_id: str, target: str, dropped: int) -> dict:
    return {
        "type": NotificationType.VULNERABILITY,
        "severity": NotificationSeverity.WARNING,
        "title": "Vulnerability Scan Incomplete",
        "message": (
            f"The scan of {target} stopped testing {dropped} "
            f"{'host' if dropped == 1 else 'hosts'} after repeated errors. "
            f"Coverage for {'that host' if dropped == 1 else 'those hosts'} is partial."
        ),
        "metadata": {
            **_scan_meta(scan_id),
            "url": f"/scans/{scan_id}?tab=vulnerabilities",
        },
    }