Mini Shell

# vim: set ts=4 sw=4 expandtab syntax=python:
"""

ngxstats.zmq_client
Realtime Nginx stats aggregation tool
ZeroMQ subscriber for fetching NGINX log data

Copyright (c) 2019-2020 InMotion Hosting, Inc.
https://www.inmotionhosting.com/

@author J. Hipps <jacobh@inmotionhosting.com>

"""

import os
import json
import pwd
import grp
import multiprocessing
import logging
import logging.handlers
from urllib.parse import urlparse

import zmq
from setproctitle import setproctitle

from ngxstats.util import gconf, check_alive

logger = logging.getLogger('ngxstats')


def start(outq):
    """
    ZeroMQ messenger
    """
    ppid = os.getppid()
    logger.info(
        "zmq_client: subservice started. pid = %d / ppid = %d",
        multiprocessing.current_process().pid,
        ppid,
    )
    try:
        setproctitle("ngxstats: zmq_client")
    except Exception:
        pass

    # Setup ZMQ context, and bind to pubsub socket
    try:
        context = zmq.Context()
        sub = context.socket(zmq.SUB)
        sub.bind(gconf.zmq_socket)
        sub.setsockopt(zmq.SUBSCRIBE, b"")
    except Exception as e:
        logger.critical(
            "zmq_client: failed to bind to socket at %s: %s",
            gconf.zmq_socket,
            str(e),
        )
        os._exit(1)  # pylint: disable=protected-access
    logger.info(
        "zmq_client: connected to pubsub socket at %s", gconf.zmq_socket
    )

    # Fix socket ownership/permissions
    spath = urlparse(gconf.zmq_socket).path
    if gconf.zmq_owner or gconf.zmq_group:
        try:
            sowner = (
                pwd.getpwnam(gconf.zmq_owner).pw_uid if gconf.zmq_owner else -1
            )
            sgroup = (
                grp.getgrnam(gconf.zmq_group).gr_gid if gconf.zmq_group else -1
            )
            os.chown(spath, sowner, sgroup)
            logger.debug(
                "zmq_client: set ownership on %s to %s:%s",
                spath,
                sowner,
                sgroup,
            )
        except Exception as e:
            logger.error(
                "zmq_client: failed to set socket ownership: %s", str(e)
            )
    if gconf.zmq_chown:
        try:
            schown = int(gconf.zmq_chown, 8)
            os.chmod(spath, schown)
            logger.debug("zmq_client: chown'd %s to %s", spath, oct(schown))
        except Exception as e:
            logger.error(
                "zmq_client: failed to set socket permissions: %s", str(e)
            )

    # Main message loop
    while True:
        # Check if master is still alive
        if not check_alive(ppid):
            logger.info("zmq_client: master %d has died. aborting.", ppid)
            os._exit(0)  # pylint: disable=protected-access

        try:
            _endpoint, msg = (
                sub.recv().decode('utf-8', errors='ignore').split('{', 1)
            )  # pylint: disable=unused-variable
            jmsg = '{' + msg
            msg_json = json.loads(jmsg)
        except Exception as e:
            logger.warning(
                "zmq_client: failed to decode message: '%s': %s", jmsg, str(e)
            )
            continue
        if gconf.zmq_debug:
            logger.debug("zmq_client: got message: %s", msg_json)
        outq.put(msg_json)