#!/usr/bin/env python3
# -*- encoding: utf-8; py-indent-offset: 4 -*-
# +------------------------------------------------------------------+
# |             ____ _               _        __  __ _  __           |
# |            / ___| |__   ___  ___| | __   |  \/  | |/ /           |
# |           | |   | '_ \ / _ \/ __| |/ /   | |\/| | ' /            |
# |           | |___| | | |  __/ (__|   <    | |  | | . \            |
# |            \____|_| |_|\___|\___|_|\_\___|_|  |_|_|\_\           |
# |                                                                  |
# | Copyright Mathias Kettner 2016             mk@mathias-kettner.de |
# +------------------------------------------------------------------+
#
# This file is part of Check_MK.
# The official homepage is at http://mathias-kettner.de/check_mk.
#
# check_mk is free software;  you can redistribute it and/or modify it
# under the  terms of the  GNU General Public License  as published by
# the Free Software Foundation in version 2.  check_mk is  distributed
# in the hope that it will be useful, but WITHOUT ANY WARRANTY;  with-
# out even the implied warranty of  MERCHANTABILITY  or  FITNESS FOR A
# PARTICULAR PURPOSE. See the  GNU General Public License for more de-
# tails. You should have  received  a copy of the  GNU  General Public
# License along with GNU Make; see the file  COPYING.  If  not,  write
# to the Free Software Foundation, Inc., 51 Franklin St,  Fifth Floor,
# Boston, MA 02110-1301 USA.

# Check_MK-Agent-Plugin - php-fpm Status
#
# By default this plugin tries to detect all locally running php-fpm processes
# and to monitor them. If this is not good for your environment you might
# create an php_fpm_pools.cfg file in MK_CONFDIR and populate the servers
# list to prevent executing the detection mechanism.

import configparser
import hashlib
import json
import os
import re
import socket
import struct
import sys
from glob import iglob as glob


# thanks to Milosz Galazka for sharing this code
# https://blog.sleeplessbeastie.eu/2019/04/01/how-to-display-php-fpm-pool-information-using-unix-socket-and-python-script/
class FCGIStatusClient:
    # FCGI protocol version
    FCGI_VERSION = 1

    # FCGI record types
    FCGI_BEGIN_REQUEST = 1
    FCGI_PARAMS = 4

    # FCGI roles
    FCGI_RESPONDER = 1

    # FCGI header length
    FCGI_HEADER_LENGTH = 8

    def __init__(self, socket_path, socket_timeout=5.0, status_path="/status", pool_name=None):
        # socket_path of type tuple is TCP
        if type(socket_path) is tuple:
            self.socket = socket.create_connection(socket_path, socket_timeout)
        else:
            self.socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
            self.socket.settimeout(socket_timeout)
            self.socket.connect(socket_path)
        self.status_path = status_path
        self.pool_name = pool_name
        self.request_id = 1

        self.params = {
            "SCRIPT_NAME": status_path,
            "SCRIPT_FILENAME": status_path,
            "QUERY_STRING": "json",
            "REQUEST_METHOD": "GET",
        }

    def close(self):
        self.socket.close()

    def define_begin_request(self):
        fcgi_begin_request = struct.pack("!HB5x", self.FCGI_RESPONDER, 0)
        fcgi_header = struct.pack(
            "!BBHHBx",
            self.FCGI_VERSION,
            self.FCGI_BEGIN_REQUEST,
            self.request_id,
            len(fcgi_begin_request),
            0,
        )
        self.fcgi_begin_request = fcgi_header + fcgi_begin_request

    def define_parameters(self):
        parameters = []
        for name, value in self.params.items():
            parameters.append(chr(len(name)) + chr(len(value)) + name + value)

        parameters = "".join(parameters)
        parameters_length = len(parameters)
        parameters_padding_req = parameters_length & 7
        parameters_padding = b"\x00" * parameters_padding_req

        fcgi_header_start = struct.pack(
            "!BBHHBx",
            self.FCGI_VERSION,
            self.FCGI_PARAMS,
            self.request_id,
            parameters_length,
            parameters_padding_req,
        )
        fcgi_header_end = struct.pack("!BBHHBx", self.FCGI_VERSION, self.FCGI_PARAMS, self.request_id, 0, 0)
        self.fcgi_params = fcgi_header_start + parameters.encode() + parameters_padding + fcgi_header_end

    def execute(self):
        self.status_data = b""
        self.socket.send(self.fcgi_begin_request)
        self.socket.send(self.fcgi_params)

        while True:
            header = self.socket.recv(self.FCGI_HEADER_LENGTH)
            fcgi_version, request_type, request_id, request_length, request_padding = struct.unpack("!BBHHBx", header)

            if request_type == 3:
                break
            elif request_type in (6, 7):
                raw_status_data = b""
                while len(raw_status_data) < request_length:
                    raw_status_data += self.socket.recv(request_length - len(raw_status_data))
                self.socket.recv(request_padding)
                self.status_data += raw_status_data

            if request_type == 7:
                raise Exception("Received an error packet: " + self.status_data.decode("ascii"))
            elif request_type not in (3, 6):
                raise Exception("Received unexpected packet type: " + str(request_type))

    def make_request(self):
        self.define_begin_request()
        self.define_parameters()
        self.execute()
        self.close()

    def get_metrics(self):
        data = json.loads(self.status_data.split(b"\r\n\r\n", 1)[1].decode("ascii"))
        fcgi_pool_name = data.pop("pool", None)
        pool_name = self.pool_name or fcgi_pool_name
        pm_type = data.pop("process manager", None)
        metrics = {"_".join(k.split()): v for k, v in data.items()}
        return pool_name, pm_type, metrics


def _print_status(pool_name, pm_type, metrics):
    for key, value in metrics.items():
        sys.stdout.write("%s %s %s %s\n" % (pool_name, pm_type, key, value))


def _parse_includes(filename, root=""):
    """
    Yield lines in filename, recursively parsing include= lines.
    root is prepended to include paths to support filesystem isolation in tests.
    """
    with open(filename, encoding="utf-8", errors="replace") as f:
        for line in f:
            if line.strip().startswith("include"):
                include = line.split("=", 1)[1].strip()
                for f in glob(root + include):
                    yield from _parse_includes(f, root=root)
            else:
                yield line


def _parse_fpm_config(f):
    """
    Parse php-fpm config, yielding dicts with status socket info
    """
    config = configparser.ConfigParser(strict=False)
    config.read_file(f)
    strip_from_listen = "'\""

    for name in config.sections():
        section = config[name]
        if "listen" not in section:
            continue

        listen = section["listen"].strip(strip_from_listen)
        if listen and not listen.startswith("/"):
            if "prefix" in section:
                listen = os.path.join(section["prefix"].strip(strip_from_listen), listen)
            else:
                # we cannot infer the relative listen anchor
                continue

        status_listen = section.get("pm.status_listen", "").strip(strip_from_listen)
        if status_listen:
            socket = status_listen
        else:
            socket = listen

        if socket:
            socket = socket.replace("$pool", name)

        yield {
            "pool_name": name,
            "socket": socket,
            "path": (section["pm.status_path"].strip(strip_from_listen) if "pm.status_path" in section else None),
        }


def _make_qualifier(configfile, taken=None):
    """
    Return a human-readable qualifier for a config file path.
    Extracts a version number (e.g. "8.1") if present and not already taken.
    Falls back to a short SHA-256 hex digest.
    """
    match = re.search(r"\d+\.\d+", configfile)
    if match:
        version = match.group(0)
        if taken is None or version not in taken:
            return version
    return hashlib.sha256(configfile.encode()).hexdigest()[:8]


def _get_worker_rss(master_pid, pool_name, proc_root="/proc"):
    """
    Return (total_rss, avg_rss) in bytes for worker processes of the given master PID and pool name,
    or (None, None) if no workers are found or the children file is missing.
    Workers are identified via the master's children list and their process title.
    """
    rss_values = []
    try:
        with open("%s/%d/task/%d/children" % (proc_root, master_pid, master_pid), encoding="utf-8") as fp:
            child_pids = [int(p) for p in fp.read().split() if p]
    except (IOError, OSError, ValueError):
        return None, None

    worker_title = "php-fpm: pool %s" % pool_name
    for pid in child_pids:
        try:
            with open("%s/%d/cmdline" % (proc_root, pid), encoding="utf-8", errors="replace") as fp:
                if fp.read().strip("\0") != worker_title:
                    continue
            with open("%s/%d/status" % (proc_root, pid), encoding="utf-8", errors="replace") as fp:
                for line in fp:
                    if line.startswith("VmRSS:"):
                        rss_values.append(int(line.split()[1]) * 1024)
                        break
        except (IOError, OSError, ValueError):
            continue

    if not rss_values:
        return None, None
    total_rss = sum(rss_values)
    avg_rss = total_rss // len(rss_values)
    return total_rss, avg_rss


def _discover_fpm(proc_root="/proc"):
    """
    Find running php-fpm processes, yielding (config file path, master PID) tuples
    """
    pattern = re.compile(r"php-fpm: master process \((.*)\)$")
    for f in glob("%s/[0-9]*/cmdline" % proc_root):
        try:
            pid = int(os.path.basename(os.path.dirname(f)))
        except ValueError:
            continue
        with open(f, encoding="utf-8", errors="replace") as fp:
            cmdline = fp.readline()
            match = pattern.match(cmdline.strip("\0"))
            if match:
                yield match.group(1), pid


def main(root="/"):
    root = root.rstrip("/")
    proc_root = root + "/proc"

    sys.stdout.write("<<<php_fpm_pools>>>\n")

    # Pass 1: collect all pools from all running php-fpm instances
    all_pools = []  # list of (configfile, fpm_status)
    master_pids = {}  # configfile -> master PID
    for configfile, master_pid in _discover_fpm(proc_root):
        for fpm_status in _parse_fpm_config(_parse_includes(root + configfile, root=root)):
            if fpm_status.get("path"):
                all_pools.append((configfile, fpm_status))
                master_pids[configfile] = master_pid

    # Pass 2: detect duplicate pool names and assign unique qualifiers
    pool_name_counts = {}
    for _, fpm_status in all_pools:
        name = fpm_status["pool_name"]
        pool_name_counts[name] = pool_name_counts.get(name, 0) + 1
    duplicate_pool_names = {name for name, count in pool_name_counts.items() if count > 1}

    used_qualifiers = {}  # pool_name -> set of qualifiers already assigned
    qualifiers = {}  # configfile -> qualifier
    for configfile, fpm_status in all_pools:
        if fpm_status["pool_name"] not in duplicate_pool_names:
            continue
        if configfile in qualifiers:
            continue
        taken = used_qualifiers.setdefault(fpm_status["pool_name"], set())
        qualifier = _make_qualifier(configfile, taken=taken)
        taken.add(qualifier)
        qualifiers[configfile] = qualifier

    # Pass 3: connect to each pool and output metrics
    for configfile, fpm_status in all_pools:
        pool_name = fpm_status["pool_name"]
        if pool_name in duplicate_pool_names:
            display_name = "%s@%s" % (pool_name, qualifiers[configfile])
        else:
            display_name = pool_name
        try:
            fcgi_client = FCGIStatusClient(
                socket_path=fpm_status.get("socket"),
                status_path=fpm_status.get("path", "/status"),
                pool_name=display_name,
            )
            fcgi_client.make_request()
            pool_name_out, pm_type, metrics = fcgi_client.get_metrics()

            total_rss, avg_rss = _get_worker_rss(master_pids[configfile], pool_name, proc_root=proc_root)
            if total_rss is not None:
                metrics["memory_total_rss"] = total_rss
                metrics["memory_avg_rss"] = avg_rss
            _print_status(pool_name_out, pm_type, metrics)

        except Exception as e:
            sys.stderr.write("Exception (%s): %s\n" % (fpm_status.get("socket"), e))


if __name__ == "__main__":  # pragma: no cover
    main()