artemis.py

from __future__ import annotations

import html
import json
import os
import re
from dataclasses import dataclass
from typing import Any

import requests
from flask import Flask, render_template_string, request

app = Flask(__name__)

# ---- Config ----
ARTEMIS_BASE_URL = os.getenv("ARTEMIS_BASE_URL", "http://localhost:8161/console/jolokia")
ARTEMIS_USER = os.getenv("ARTEMIS_USER", "***")
ARTEMIS_PASSWORD = os.getenv("ARTEMIS_PASSWORD", "***")
ARTEMIS_BROKER_NAME = os.getenv("ARTEMIS_BROKER_NAME", "default")
REQUEST_TIMEOUT = int(os.getenv("ARTEMIS_TIMEOUT", "15"))
DEFAULT_SHOW_ALL = os.getenv("SHOW_ALL_QUEUES_BY_DEFAULT", "false").lower() == "true"

session = requests.Session()
session.auth = (ARTEMIS_USER, ARTEMIS_PASSWORD)
session.headers.update({"Content-Type": "application/json"})


@dataclass
class QueueRef:
    name: str
    address: str
    routing_type: str
    mbean: str


@dataclass
class QueueSummary:
    name: str
    count: int
    sample_messages: list[dict[str, Any]]
    error: str | None = None
    address: str | None = None
    routing_type: str | None = None


@dataclass
class QueueCountRow:
    name: str
    address: str
    routing_type: str
    message_count: int | None
    error: str | None = None


@dataclass
class DlqStats:
    queue_count: int
    total_messages: int


def jolokia_request(payload: dict[str, Any]) -> Any:
    response = session.post(ARTEMIS_BASE_URL, json=payload, timeout=REQUEST_TIMEOUT)
    response.raise_for_status()
    data = response.json()

    if data.get("status") != 200:
        raise RuntimeError(f"Jolokia call failed: {data}")

    return data.get("value")


def jolokia_read(mbean: str, attribute: str) -> Any:
    return jolokia_request(
        {
            "type": "read",
            "mbean": mbean,
            "attribute": attribute,
        }
    )


def jolokia_exec(mbean: str, operation: str, arguments: list[Any] | None = None) -> Any:
    return jolokia_request(
        {
            "type": "exec",
            "mbean": mbean,
            "operation": operation,
            "arguments": arguments or [],
        }
    )


def jolokia_search(mbean_pattern: str) -> list[str]:
    value = jolokia_request(
        {
            "type": "search",
            "mbean": mbean_pattern,
        }
    )
    if not value:
        return []
    return list(value)


def normalise_jmx_value(value: str) -> str:
    return value.replace(r"\.", ".").replace(r"\"", '"').replace(r"\\", "\\")


def extract_mbean_property(mbean: str, property_name: str) -> str | None:
    pattern = rf'{re.escape(property_name)}="((?:\\.|[^"])*)"'
    match = re.search(pattern, mbean)
    if not match:
        return None
    return normalise_jmx_value(match.group(1))


def get_all_queue_refs() -> list[QueueRef]:
    pattern = (
        f'org.apache.activemq.artemis:broker="{ARTEMIS_BROKER_NAME}",'
        f'component=addresses,address=*,subcomponent=queues,routing-type=*,queue=*'
    )

    mbeans = jolokia_search(pattern)
    queue_refs: list[QueueRef] = []

    for mbean in mbeans:
        queue_name = extract_mbean_property(mbean, "queue")
        address = extract_mbean_property(mbean, "address")
        routing_type = extract_mbean_property(mbean, "routing-type")

        if not queue_name:
            continue

        queue_refs.append(
            QueueRef(
                name=queue_name,
                address=address or queue_name,
                routing_type=routing_type or "unknown",
                mbean=mbean,
            )
        )

    queue_refs.sort(key=lambda q: (q.name.lower(), q.routing_type.lower(), q.address.lower()))
    return queue_refs


def is_dlq(queue_ref: QueueRef) -> bool:
    text = f"{queue_ref.name} {queue_ref.address}".lower()
    return "dlq" in text or "dead" in text


def list_messages_for_queue_mbean(mbean: str, selector: str = "") -> list[dict[str, Any]]:
    raw = jolokia_exec(mbean, "listMessages(java.lang.String)", [selector])

    if raw is None:
        return []

    if isinstance(raw, list):
        return [m if isinstance(m, dict) else {"raw": m} for m in raw]

    if isinstance(raw, dict):
        return [raw]

    if isinstance(raw, str):
        try:
            parsed = json.loads(raw)
            if isinstance(parsed, list):
                return [m if isinstance(m, dict) else {"raw": m} for m in parsed]
            if isinstance(parsed, dict):
                return [parsed]
            return [{"raw": parsed}]
        except json.JSONDecodeError:
            return [{"raw": raw}]

    return [{"raw": raw}]


def get_message_count_for_queue_mbean(mbean: str) -> int:
    value = jolokia_read(mbean, "MessageCount")
    if value is None:
        return 0
    return int(value)


def get_dlq_stats(queue_refs: list[QueueRef]) -> DlqStats:
    dlq_refs = [q for q in queue_refs if is_dlq(q)]
    total_messages = 0

    for q in dlq_refs:
        try:
            total_messages += get_message_count_for_queue_mbean(q.mbean)
        except Exception:
            # Keep the page usable even if one queue metric fails.
            pass

    return DlqStats(
        queue_count=len(dlq_refs),
        total_messages=total_messages,
    )


def safe_contains(message: dict[str, Any], text: str) -> bool:
    haystack = json.dumps(message, default=str).lower()
    return text.lower() in haystack


BASE_STYLES = """
<style>
    body { font-family: Arial, sans-serif; margin: 24px; background: #fafafa; color: #222; }
    h1 { margin-bottom: 8px; }
    form { margin-bottom: 20px; padding: 12px; background: white; border: 1px solid #ddd; border-radius: 6px; }
    input[type=text] { width: 360px; padding: 8px; }
    button { padding: 8px 12px; }
    label { margin-right: 12px; }
    .summary { margin: 16px 0; }
    .queue { background: white; border: 1px solid #ddd; margin-bottom: 16px; padding: 12px; border-radius: 6px; }
    .queue h2 { margin: 0 0 8px 0; font-size: 18px; }
    .meta { color: #666; font-size: 13px; margin-bottom: 8px; }
    .count { color: #555; font-size: 14px; margin-bottom: 8px; }
    .error { color: #b00020; white-space: pre-wrap; word-wrap: break-word; }
    details { margin-top: 8px; }
    pre { white-space: pre-wrap; word-wrap: break-word; background: #f5f5f5; padding: 10px; border-radius: 4px; }
    .muted { color: #666; }
    .topnote { margin-bottom: 12px; color: #555; }
    .nav { margin-bottom: 16px; }
    .nav a { margin-right: 12px; }
    table { border-collapse: collapse; width: 100%; background: white; }
    th, td { text-align: left; padding: 10px; border-bottom: 1px solid #ddd; }
    th { background: #f3f3f3; }
    .cards { display: grid; grid-template-columns: repeat(auto-fit, minmax(220px, 1fr)); gap: 12px; margin: 0 0 18px 0; }
    .card { background: white; border: 1px solid #ddd; border-radius: 6px; padding: 14px; }
    .card-title { color: #666; font-size: 13px; margin-bottom: 6px; }
    .card-value { font-size: 28px; font-weight: 700; }
</style>
"""

SEARCH_TEMPLATE = """
<!doctype html>
<html lang="en">
<head>
    <meta charset="utf-8">
    <title>Artemis Queue Aggregator</title>
    {{ base_styles|safe }}
</head>
<body>
    <div class="nav">
        <a href="/">Search</a>
        <a href="/queues">Queue counts</a>
    </div>

    <h1>Artemis Queue Aggregator</h1>
    <div class="topnote">Broker: {{ broker_name }}</div>

    <div class="cards">
        <div class="card">
            <div class="card-title">DLQ queues</div>
            <div class="card-value">{{ dlq_stats.queue_count }}</div>
        </div>
        <div class="card">
            <div class="card-title">Total DLQ messages</div>
            <div class="card-value">{{ dlq_stats.total_messages }}</div>
        </div>
    </div>

    <form method="get">
        <label for="q">Search text:</label>
        <input id="q" name="q" type="text" value="{{ search_text }}" placeholder="caseId, hearingId, event name...">
        <label>
            <input type="checkbox" name="all" value="1" {% if show_all %}checked{% endif %}>
            Include all queues
        </label>
        <button type="submit">Search</button>
    </form>

    <div class="summary">
        <strong>{{ queue_summaries|length }}</strong> queues shown
        {% if not search_text %}
            <span class="muted">— queue names only; enter a search to query messages</span>
        {% endif %}
    </div>

    {% for queue in queue_summaries %}
        <div class="queue">
            <h2>{{ queue.name }}</h2>
            <div class="meta">
                address={{ queue.address }} | routing-type={{ queue.routing_type }}
            </div>
            {% if queue.error %}
                <div class="error">{{ queue.error }}</div>
            {% else %}
                <div class="count">{{ queue.count }} matching messages</div>

                {% if queue.sample_messages %}
                    {% for message in queue.sample_messages %}
                        <details>
                            <summary>Message {{ loop.index }}</summary>
                            <pre>{{ message | tojson(indent=2) }}</pre>
                        </details>
                    {% endfor %}
                {% elif search_text %}
                    <div class="muted">No matching messages</div>
                {% endif %}
            {% endif %}
        </div>
    {% endfor %}
</body>
</html>
"""

QUEUE_COUNTS_TEMPLATE = """
<!doctype html>
<html lang="en">
<head>
    <meta charset="utf-8">
    <title>Artemis Queue Counts</title>
    {{ base_styles|safe }}
</head>
<body>
    <div class="nav">
        <a href="/">Search</a>
        <a href="/queues">Queue counts</a>
    </div>

    <h1>Artemis Queue Counts</h1>
    <div class="topnote">Broker: {{ broker_name }}</div>

    <div class="cards">
        <div class="card">
            <div class="card-title">DLQ queues</div>
            <div class="card-value">{{ dlq_stats.queue_count }}</div>
        </div>
        <div class="card">
            <div class="card-title">Total DLQ messages</div>
            <div class="card-value">{{ dlq_stats.total_messages }}</div>
        </div>
    </div>

    <form method="get">
        <label for="filter">Filter queue name:</label>
        <input id="filter" name="filter" type="text" value="{{ filter_text }}" placeholder="DLQ, hearing, subscriptions...">
        <label>
            <input type="checkbox" name="all" value="1" {% if show_all %}checked{% endif %}>
            Include all queues
        </label>
        <button type="submit">Refresh</button>
    </form>

    <div class="summary">
        <strong>{{ rows|length }}</strong> queues shown
    </div>

    <table>
        <thead>
            <tr>
                <th>Queue</th>
                <th>Address</th>
                <th>Routing Type</th>
                <th>Message Count</th>
            </tr>
        </thead>
        <tbody>
            {% for row in rows %}
                <tr>
                    <td>{{ row.name }}</td>
                    <td>{{ row.address }}</td>
                    <td>{{ row.routing_type }}</td>
                    <td>
                        {% if row.error %}
                            <span class="error">{{ row.error }}</span>
                        {% else %}
                            {{ row.message_count }}
                        {% endif %}
                    </td>
                </tr>
            {% endfor %}
        </tbody>
    </table>
</body>
</html>
"""


@app.route("/")
def index():
    search_text = request.args.get("q", "").strip()
    show_all = request.args.get("all") == "1" or DEFAULT_SHOW_ALL
    queue_summaries: list[QueueSummary] = []

    try:
        queue_refs = get_all_queue_refs()
        dlq_stats = get_dlq_stats(queue_refs)
    except Exception as exc:
        return f"<h1>Failed to load queue MBeans</h1><pre>{html.escape(str(exc))}</pre>", 500

    if not show_all:
        queue_refs = [q for q in queue_refs if is_dlq(q)]

    for queue_ref in queue_refs:
        try:
            if search_text:
                messages = list_messages_for_queue_mbean(queue_ref.mbean)
                messages = [m for m in messages if safe_contains(m, search_text)]
                count = len(messages)
                sample_messages = messages[:3]
            else:
                count = 0
                sample_messages = []

            queue_summaries.append(
                QueueSummary(
                    name=queue_ref.name,
                    count=count,
                    sample_messages=sample_messages,
                    address=queue_ref.address,
                    routing_type=queue_ref.routing_type,
                )
            )
        except Exception as exc:
            queue_summaries.append(
                QueueSummary(
                    name=queue_ref.name,
                    count=0,
                    sample_messages=[],
                    error=str(exc),
                    address=queue_ref.address,
                    routing_type=queue_ref.routing_type,
                )
            )

    if search_text:
        queue_summaries = [q for q in queue_summaries if q.count > 0 or q.error]

    return render_template_string(
        SEARCH_TEMPLATE,
        queue_summaries=queue_summaries,
        search_text=search_text,
        broker_name=ARTEMIS_BROKER_NAME,
        show_all=show_all,
        base_styles=BASE_STYLES,
        dlq_stats=dlq_stats,
    )


@app.route("/queues")
def queues():
    filter_text = request.args.get("filter", "").strip().lower()
    show_all = True if request.args.get("all") is None else request.args.get("all") == "1"
    rows: list[QueueCountRow] = []

    try:
        queue_refs = get_all_queue_refs()
        dlq_stats = get_dlq_stats(queue_refs)
    except Exception as exc:
        return f"<h1>Failed to load queue MBeans</h1><pre>{html.escape(str(exc))}</pre>", 500

    if not show_all:
        queue_refs = [q for q in queue_refs if is_dlq(q)]

    if filter_text:
        queue_refs = [
            q for q in queue_refs
            if filter_text in q.name.lower() or filter_text in q.address.lower()
        ]

    for queue_ref in queue_refs:
        try:
            count = get_message_count_for_queue_mbean(queue_ref.mbean)
            rows.append(
                QueueCountRow(
                    name=queue_ref.name,
                    address=queue_ref.address,
                    routing_type=queue_ref.routing_type,
                    message_count=count,
                )
            )
        except Exception as exc:
            rows.append(
                QueueCountRow(
                    name=queue_ref.name,
                    address=queue_ref.address,
                    routing_type=queue_ref.routing_type,
                    message_count=None,
                    error=str(exc),
                )
            )

    rows.sort(key=lambda r: (-1 if r.message_count is None else -r.message_count, r.name.lower()))

    return render_template_string(
        QUEUE_COUNTS_TEMPLATE,
        rows=rows,
        filter_text=filter_text,
        broker_name=ARTEMIS_BROKER_NAME,
        show_all=show_all,
        base_styles=BASE_STYLES,
        dlq_stats=dlq_stats,
    )


if __name__ == "__main__":
    app.run(debug=True, host="0.0.0.0", port=8081)