BotZ/bot_z/cli.py

# -*- coding: utf-8 -*-

"""Console script to control the bot_z daemon"""

import errno
import logging
import os
import sys
import time
import typing as T

import click
import lockfile

from bot_z.zdaemon import daemon_process, Fifo

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
sh = logging.StreamHandler(stream=sys.stdout)
sh.setFormatter(logging.Formatter("%(message)s"))
logger.addHandler(sh)


@click.group()
@click.option("-d", "--debug", is_flag=True, default=False, help="Enable debug mode.")
@click.option(
    "-v",
    "--verbose",
    is_flag=True,
    default=False,
    help="More verbose output from this cli.",
)
@click.option(
    "-f",
    "--fifo",
    type=click.Path(),
    default="/tmp/bot_z.cmd",
    help="Path to the control fifo.",
)
@click.option(
    "-w",
    "--workdir",
    type=click.Path(exists=True, readable=True, writable=True, resolve_path=True),
    default="/tmp/",
    help="The working dir where to launch the daemon and where the lockfile is put.",
)
@click.pass_context
def cli(
    ctx: click.Context, debug: bool, verbose: bool, fifo: str, workdir: str
) -> None:
    """
    Group cli.
    """
    if verbose:
        logger.setLevel(logging.DEBUG)
    ctx.ensure_object(dict)
    ctx.obj["debug"] = debug
    ctx.obj["verbose"] = verbose
    ctx.obj["fifo"] = fifo
    ctx.obj["workdir"] = workdir


@cli.command("start")
@click.option(
    "-t", "--timeout", type=click.INT, default=20, help="Browser requests timeout."
)
@click.option(
    "-m",
    "--umask",
    type=click.INT,
    default=0o002,
    help="The umask of the control fifo.",
)
@click.option(
    "-n",
    "--name",
    type=click.STRING,
    default="default_instance",
    help="The daemon instance name.",
)
@click.option(
    "-p",
    "--proxy",
    type=click.STRING,
    default=None,
    help="An optional string for the proxy with the form 'address:port'.",
)
@click.option(
    "--foreground",
    is_flag=True,
    default=False,
    help="Keep the process in foreground (do not daemonize).",
)
@click.argument("baseuri", type=click.STRING)
@click.pass_context
def start_command(
    ctx: click.Context,
    baseuri: str,
    name: str,
    umask: int,
    timeout: int,
    proxy: T.Optional[str] = None,
    foreground: bool = False,
) -> None:
    """
    Invokes daemon_process for the first time.
    """
    lf = lockfile.FileLock(os.path.join(ctx.obj["workdir"], "bot_z"))
    logger.debug("Daemon starting. Lockfile: %s", lf)
    proxy_tuple = None
    if proxy:
        proxy_tuple = tuple(proxy.split(":"))

    daemon_process(
        fifo_path=ctx.obj["fifo"],
        working_dir=ctx.obj["workdir"],
        umask=umask,
        pidfile=lf,
        base_uri=baseuri,
        name=name,
        timeout=timeout,
        proxy=proxy_tuple,
        headless=not ctx.obj["debug"],
        debug=ctx.obj["debug"],
        foreground=foreground,
    )
    logger.info("Daemon started.")


@cli.command("stop")
@click.pass_context
def stop_command(ctx: click.Context):
    """
    Writes on the fifo. Invokes the stop.
    """
    logger.debug("Sending the stop command down the pipe: %s", ctx.obj["fifo"])
    with open(ctx.obj["fifo"], "w") as fifo:
        fifo.write("stop")
        fifo.flush()
    logger.info("Stop sent.")


@cli.command("reload")
@click.pass_context
def reload_command(ctx: click.Context):
    """
    Writes on the fifo. Invokes the reload.
    """
    logger.debug("Sending the reload command down the pipe: %s", ctx.obj["fifo"])
    with open(ctx.obj["fifo"], "w") as fifo:
        fifo.write("reload")
        fifo.flush()
    logger.info("Reload sent.")


@cli.command("status")
@click.pass_context
def status_command(ctx: click.Context):
    """
    Writes on the fifo. Queries the status.
    """
    rfifo_path = os.path.join(ctx.obj["workdir"], "rfifo-{}".format(id(ctx)))
    logger.debug("Awaiting response on fifo: %s", rfifo_path)
    try:
        os.mkfifo(rfifo_path)
        logger.debug("Response fifo newly created.")
    except OSError as err:
        if err.errno != errno.EEXIST:
            raise
        logger.debug("Response fifo exists.")

    with Fifo(ctx.obj["fifo"], "w") as fifo:
        fifo.write("status|{}".format(rfifo_path))
    logger.debug("Awaiting response...")
    done = False
    while not done:
        try:
            with open(rfifo_path, "r") as rfifo:
                resp = rfifo.read()
            done = True
        except FileNotFoundError as e:
            logger.warning("File not found: %s", e)
            pass
    logger.info(resp)
    try:
        os.remove(rfifo_path)
    except OSError as e:
        logger.warning("Failed removing response fifo %s: %s", rfifo_path, e)
        pass


if __name__ == "__main__":
    cli()