Nagios-compatible health check for Amazon MSK

check_msk_produce_consume.py

#!/usr/bin/env python3.11
#
# Nagios-compatible health check for Amazon MSK.
#
# The script performs an end-to-end Kafka check against an MSK cluster using IAM authentication:
#
# 1. Fetches topic metadata.
# 2. Selects a topic partition.
# 3. Reads the current end offset.
# 4. Produces a unique test message to that exact partition.
# 5. Consumes from the saved offset.
# 6. Verifies that the produced message is received.
#
# It avoids Kafka CLI tools and uses `confluent-kafka` directly, so it is suitable for Nagios/NRPE checks.
# The output includes standard Nagios status codes and perfdata for graphing latency metrics such as metadata,
# offset lookup, produce, consume, and total time.
#
# Requires:
#   pip3.11 install confluent-kafka aws-msk-iam-sasl-signer-python
#
# nrpe command:
#   command[check_msk_produce_consume]=/usr/local/bin/check_msk_produce_consume.py --bootstrap $ARG1$ --topic nagios-msk-healthcheck --auth iam --region us-east-1 --group nagios-msk-healthcheck --warning 5 --critical 10
#
# Required IAM policy:
# {
#     "Version": "2012-10-17",
#     "Statement": [
#         {
#             "Sid": "VisualEditor0",
#             "Effect": "Allow",
#             "Action": [
#                 "kafka-cluster:*Topic*",
#                 "kafka-cluster:DescribeCluster",
#                 "kafka-cluster:ReadData",
#                 "kafka-cluster:Connect",
#                 "kafka-cluster:WriteData"
#             ],
#             "Resource": [
#                 "arn:aws:kafka:us-east-1:account id:cluster/cluster name/*",
#                 "arn:aws:kafka:*:account id:topic/cluster name/*/nagios-msk-healthcheck"
#             ]
#         },
#         {
#             "Effect": "Allow",
#             "Action": [
#                 "kafka-cluster:DescribeGroup",
#                 "kafka-cluster:AlterGroup"
#             ],
#             "Resource": "arn:aws:kafka:us-east-1:account id:group/cluster name/*/nagios-msk-healthcheck"
#         }
#     ]
# }
#
# https://gist.github.com/jazzl0ver/30cb5118fb0874465211275b02692f89
#

import argparse
import os
import random
import sys
import time
import uuid

from confluent_kafka import Consumer, Producer, TopicPartition, KafkaException
from confluent_kafka.admin import AdminClient
from aws_msk_iam_sasl_signer import MSKAuthTokenProvider


OK = 0
WARNING = 1
CRITICAL = 2
UNKNOWN = 3


def nagios_exit(code, msg, perfdata=None):
    print(f"{msg} | {perfdata}" if perfdata else msg)
    sys.exit(code)


def oauth_cb(region):
    def _oauth_cb(oauth_config):
        token, expiry_ms = MSKAuthTokenProvider.generate_auth_token(region)
        return token, expiry_ms / 1000

    return _oauth_cb


def common_kafka_config(args, client_id):
    cfg = {
        "bootstrap.servers": args.bootstrap,
        "client.id": client_id,
        "socket.timeout.ms": args.socket_timeout_ms,
        "log.connection.close": False,
        # suppress librdkafka telemetry/info noise
        "enable.metrics.push": False,
        "log_level": 3,
    }

    if args.auth == "iam":
        cfg.update({
            "security.protocol": "SASL_SSL",
            "sasl.mechanisms": "OAUTHBEARER",
            "oauth_cb": oauth_cb(args.region),
        })
        return cfg

    if args.auth == "ssl":
        cfg.update({
            "security.protocol": "SSL",
        })

        if args.ca_location:
            cfg["ssl.ca.location"] = args.ca_location

        return cfg

    nagios_exit(UNKNOWN, f"UNKNOWN - unsupported auth mode: {args.auth}")


def admin_config(args):
    cfg = common_kafka_config(args, "nagios-msk-admin")
    cfg.update({
        "request.timeout.ms": args.request_timeout_ms,
    })
    return cfg


def producer_config(args):
    cfg = common_kafka_config(args, "nagios-msk-producer")
    cfg.update({
        "request.timeout.ms": args.request_timeout_ms,
        "message.timeout.ms": args.message_timeout_ms,
        "acks": "all",
    })
    return cfg


def consumer_config(args, client_id):
    cfg = common_kafka_config(args, client_id)
    cfg.update({
        "group.id": args.group,
        "enable.auto.commit": False,
        "auto.offset.reset": "error",
    })
    return cfg


def get_partitions(args):
    admin = AdminClient(admin_config(args))

    start = time.time()
    md = admin.list_topics(topic=args.topic, timeout=args.timeout)
    elapsed = time.time() - start

    if args.topic not in md.topics:
        nagios_exit(CRITICAL, f"CRITICAL - topic not found: {args.topic}")

    topic_md = md.topics[args.topic]

    if topic_md.error is not None:
        nagios_exit(CRITICAL, f"CRITICAL - topic metadata error: {topic_md.error}")

    partitions = sorted(topic_md.partitions.keys())

    if not partitions:
        nagios_exit(CRITICAL, f"CRITICAL - topic has no partitions: {args.topic}")

    return partitions, elapsed


def choose_partition(args, partitions):
    if args.partition is not None:
        if args.partition not in partitions:
            nagios_exit(
                CRITICAL,
                (
                    f"CRITICAL - partition {args.partition} not found in topic "
                    f"{args.topic}; available={partitions}"
                ),
            )
        return args.partition

    return random.choice(partitions)


def get_end_offset(args, partition):
    consumer = Consumer(consumer_config(args, "nagios-msk-offset-check"))
    tp = TopicPartition(args.topic, partition)

    start = time.time()

    try:
        low, high = consumer.get_watermark_offsets(tp, timeout=args.timeout)
        return high, time.time() - start
    finally:
        consumer.close()


def produce_message(args, partition, key, value):
    producer = Producer(producer_config(args))

    result = {
        "error": None,
        "delivered": False,
        "offset": None,
    }

    def delivery_cb(err, msg):
        if err is not None:
            result["error"] = err
            return

        result["delivered"] = True
        result["offset"] = msg.offset()

    start = time.time()

    try:
        producer.produce(
            topic=args.topic,
            key=key.encode("utf-8"),
            value=value.encode("utf-8"),
            partition=partition,
            callback=delivery_cb,
        )
        producer.flush(args.timeout)

    except BufferError as e:
        nagios_exit(CRITICAL, f"CRITICAL - producer buffer error: {e}")

    except KafkaException as e:
        nagios_exit(CRITICAL, f"CRITICAL - producer Kafka error: {e}")

    elapsed = time.time() - start

    if result["error"] is not None:
        nagios_exit(CRITICAL, f"CRITICAL - produce failed: {result['error']}")

    if not result["delivered"]:
        nagios_exit(CRITICAL, "CRITICAL - produce failed: delivery callback not confirmed")

    return elapsed, result["offset"]


def consume_message(args, partition, offset, expected_key, expected_value):
    consumer = Consumer(consumer_config(args, "nagios-msk-consumer"))
    tp = TopicPartition(args.topic, partition, offset)

    start = time.time()
    deadline = start + args.consume_timeout

    try:
        consumer.assign([tp])

        while time.time() < deadline:
            msg = consumer.poll(1.0)

            if msg is None:
                continue

            if msg.error():
                nagios_exit(CRITICAL, f"CRITICAL - consume error: {msg.error()}")

            key = msg.key().decode("utf-8") if msg.key() else ""
            value = msg.value().decode("utf-8") if msg.value() else ""

            if key == expected_key and value == expected_value:
                return time.time() - start, msg.offset()

            if args.verbose:
                print(
                    (
                        f"DEBUG - skipped message partition={msg.partition()} "
                        f"offset={msg.offset()} key={key} value={value}"
                    ),
                    file=sys.stderr,
                )

        nagios_exit(
            CRITICAL,
            (
                f"CRITICAL - message not found in topic={args.topic} "
                f"partition={partition} from_offset={offset}"
            ),
        )

    finally:
        consumer.close()


def main():
    parser = argparse.ArgumentParser(
        description="Nagios check for Amazon MSK produce/consume using partition+offset"
    )

    parser.add_argument(
        "-b", "--bootstrap",
        required=True,
        help="Bootstrap servers, e.g. b-1.xxx:9098,b-2.xxx:9098",
    )
    parser.add_argument(
        "-t", "--topic",
        required=True,
        help="Topic name",
    )
    parser.add_argument(
        "--auth",
        choices=["iam", "ssl"],
        default="iam",
        help="Auth mode. Default: iam",
    )
    parser.add_argument(
        "--region",
        default=os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION"),
        help="AWS region for MSK IAM, e.g. us-east-1",
    )
    parser.add_argument(
        "--group",
        default="nagios-msk-healthcheck",
        help="Consumer group id required by confluent-kafka",
    )
    parser.add_argument(
        "--partition",
        type=int,
        default=None,
        help="Partition to use. Default: random partition",
    )
    parser.add_argument(
        "--ca-location",
        help="CA bundle path for SSL mode",
    )
    parser.add_argument(
        "--timeout",
        type=float,
        default=10.0,
        help="General timeout seconds",
    )
    parser.add_argument(
        "--consume-timeout",
        type=float,
        default=15.0,
        help="Consume timeout seconds",
    )
    parser.add_argument(
        "--socket-timeout-ms",
        type=int,
        default=10000,
        help="Kafka socket timeout ms",
    )
    parser.add_argument(
        "--request-timeout-ms",
        type=int,
        default=10000,
        help="Kafka request timeout ms",
    )
    parser.add_argument(
        "--message-timeout-ms",
        type=int,
        default=10000,
        help="Producer message timeout ms",
    )
    parser.add_argument(
        "-w", "--warning",
        type=float,
        default=5.0,
        help="Warning threshold for total seconds",
    )
    parser.add_argument(
        "-c", "--critical",
        type=float,
        default=10.0,
        help="Critical threshold for total seconds",
    )
    parser.add_argument(
        "-v", "--verbose",
        action="store_true",
        help="Verbose debug to stderr",
    )

    args = parser.parse_args()

    if args.auth == "iam" and not args.region:
        nagios_exit(UNKNOWN, "UNKNOWN - --region is required for IAM auth")

    total_start = time.time()

    key = f"nagios-msk-key-{uuid.uuid4()}"
    value = f"nagios-msk-value-{uuid.uuid4()}"

    try:
        partitions, metadata_time = get_partitions(args)
        partition = choose_partition(args, partitions)

        start_offset, offset_time = get_end_offset(args, partition)

        produce_time, produced_offset = produce_message(
            args=args,
            partition=partition,
            key=key,
            value=value,
        )

        consume_time, consumed_offset = consume_message(
            args=args,
            partition=partition,
            offset=start_offset,
            expected_key=key,
            expected_value=value,
        )

    except KafkaException as e:
        nagios_exit(CRITICAL, f"CRITICAL - Kafka error: {e}")

    except Exception as e:
        nagios_exit(UNKNOWN, f"UNKNOWN - unexpected error: {e}")

    total_time = time.time() - total_start
    check_time = produce_time + consume_time

    perfdata = (
        f"check_time={check_time:.3f}s;{args.warning:g};{args.critical:g};0; "
        f"total_time={total_time:.3f}s;;;0; "
        f"metadata_time={metadata_time:.3f}s;;;0; "
        f"offset_time={offset_time:.3f}s;;;0; "
        f"produce_time={produce_time:.3f}s;;;0; "
        f"consume_time={consume_time:.3f}s;;;0;"
    )

    details = (
        f"published in {produce_time:.3f}, consumed in {consume_time:.3f}, check_time={check_time:.3f}s; total={total_time:.3f}s"
    )

    if check_time >= args.critical:
        nagios_exit(
            CRITICAL,
            f"CRITICAL - produce/consume OK but slow: {details}",
            perfdata,
        )

    if check_time >= args.warning:
        nagios_exit(
            WARNING,
            f"WARNING - produce/consume OK but slow: {details}",
            perfdata,
        )

    nagios_exit(
        OK,
        f"OK - produce/consume OK: {details}",
        perfdata,
    )


if __name__ == "__main__":
    main()