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 typing as T

import click
import lockfile

from bot_z.zdaemon import daemon_process
from bot_z.utils import Fifo, PLifo, cmd_marshal

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


def _check_name(lifo_path: str, name: T.Optional[str]) -> str:
    with PLifo(lifo_path) as lifo:
        if name is None:
            return lifo.last()
        if name not in lifo:
            logger.error("Daemon not present. Exiting.")
            sys.exit(2)
        return name


@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
    ctx.obj["lifo"] = os.path.join(workdir, "botz_open_daemons.list")


@cli.command("start-daemon")
@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(
    "-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_daemon_command(
    ctx: click.Context,
    baseuri: 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,
        timeout=timeout,
        proxy=proxy_tuple,
        headless=not ctx.obj["debug"],
        debug=ctx.obj["debug"],
        foreground=foreground,
    )
    logger.info("Daemon started.")


@cli.command("list")
@click.option(
    "-s",
    "--status",
    is_flag=True,
    default=False,
    help="Show also the status of each client.",
)
@click.pass_context
def list_command(ctx: click.Context, status: bool) -> None:
    """
    Shows the clients attached
    """
    logger.debug('Invoked the "list" command.')
    with PLifo(ctx.obj["lifo"]) as lifo:
        if len(lifo) == 0:
            logger.info("No clients.")
            return
        for client in lifo.all():
            logger.info(client)
            if status:
                _status_command(ctx, client)


@cli.command("start")
@click.option(
    "-n",
    "--name",
    type=click.STRING,
    prompt="Give the instance a name",
    help="The daemon instance name.",
)
@click.pass_context
def start_command(ctx: click.Context, name: str) -> None:
    """
    Writes on the fifo. Invokes the start of a client.
    """
    logger.debug("Sending the start command down the pipe: %s", ctx.obj["fifo"])
    with open(ctx.obj["fifo"], "w") as fifo:
        fifo.write(cmd_marshal(name=name, cmd="start"))
    logger.info("Start sent.")


@cli.command("stop")
@click.option("-n", "--name", default=None, help="The instance to interact with.")
@click.pass_context
def stop_command(ctx: click.Context, name: T.Optional[str]) -> None:
    """
    Writes on the fifo. Invokes the stop.
    """
    logger.debug("Sending the stop command down the pipe: %s", ctx.obj["fifo"])
    name = _check_name(ctx.obj["lifo"], name)
    logging.info("Stopping instance: %s", name)
    with open(ctx.obj["fifo"], "w") as fifo:
        fifo.write(cmd_marshal(name=name, cmd="stop"))
        fifo.flush()
    logger.info("Stop sent.")


@cli.command("reload")
@click.option("-n", "--name", default=None, help="The instance to interact with.")
@click.pass_context
def reload_command(ctx: click.Context, name: T.Optional[str]) -> None:
    """
    Writes on the fifo. Invokes the reload.
    """
    logger.debug("Sending the reload command down the pipe: %s", ctx.obj["fifo"])
    name = _check_name(ctx.obj["lifo"], name)
    with open(ctx.obj["fifo"], "w") as fifo:
        fifo.write(cmd_marshal(name=name, cmd="reload"))
        fifo.flush()
    logger.info("Reload sent.")


@cli.command("status")
@click.option("-n", "--name", default=None, help="The instance to interact with.")
@click.pass_context
def status_command(ctx: click.Context, name: T.Optional[str]) -> None:
    """
    Writes on the fifo. Queries the status.
    """
    name = _check_name(ctx.obj["lifo"], name)
    _status_command(ctx, name)


def _status_command(ctx: click.Context, name: str) -> None:
    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(cmd_marshal(name=name, cmd="status", rfifo_path=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()