diff --git a/bot_z/bot_z.py b/bot_z/bot_z.py index dc8b213..76561a2 100644 --- a/bot_z/bot_z.py +++ b/bot_z/bot_z.py @@ -43,6 +43,7 @@ logging.basicConfig( m_logger = logging.getLogger(__name__) m_logger.debug("Init at debug") + def safely(retries: int = 0) -> T.Callable: def _safely(f: T.Callable) -> T.Callable: def _protection(self, *args, **kwargs): diff --git a/bot_z/cli.py b/bot_z/cli.py index 5cedd27..e80304b 100644 --- a/bot_z/cli.py +++ b/bot_z/cli.py @@ -1,16 +1,158 @@ # -*- coding: utf-8 -*- -"""Console script for bot_z.""" +"""Console script to control the bot_z daemon""" + +import errno +import os +import time +import typing as T import click +import lockfile + +from bot_z.zdaemon import daemon_process, Fifo -@click.command() -def main(args=None): - """Console script for bot_z.""" - click.echo("Replace this message by putting your code into " "bot_z.cli.main") - click.echo("See click documentation at http://click.pocoo.org/") +@click.group() +@click.option("-d", "--debug", is_flag=True, default=False, help="Enable debug mode.") +@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, fifo: str, workdir: str) -> None: + """ + Group cli. + """ + ctx.ensure_object(dict) + ctx.obj["debug"] = debug + 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.lock")) + 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, + ) + + +@cli.command("stop") +@click.pass_context +def stop_command(ctx: click.Context): + """ + Writes on the fifo. Invokes the stop. + """ + with open(ctx.obj["fifo"], "w") as fifo: + fifo.write("stop") + fifo.flush() + + +@cli.command("reload") +@click.pass_context +def reload_command(ctx: click.Context): + """ + Writes on the fifo. Invokes the reload. + """ + with open(ctx.obj["fifo"], "w") as fifo: + fifo.write("reload") + fifo.flush() + + +@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))) + try: + os.mkfifo(rfifo_path) + except OSError as err: + if err.errno != errno.EEXIST: + raise + + with Fifo(ctx.obj["fifo"], "w") as fifo: + fifo.write("status|{}".format(rfifo_path)) + done = False + while not done: + try: + with open(rfifo_path, "r") as rfifo: + resp = rfifo.read() + done = True + except FileNotFoundError as e: + print(e) + pass + print(resp) if __name__ == "__main__": - main() + cli() diff --git a/bot_z/zdaemon.py b/bot_z/zdaemon.py new file mode 100644 index 0000000..d72c0be --- /dev/null +++ b/bot_z/zdaemon.py @@ -0,0 +1,195 @@ +# -*- coding: utf-8 -*- + +"""Daemon for local execution of commands.""" + +import errno +import functools +import json +import lockfile +import logging +import logging.handlers +import os +import signal +import typing as T + +from bot_z.bot_z import Operator +import daemon +from selenium.webdriver.common.keys import Keys + + +logging.basicConfig( + level=os.environ.get("BOTZ_LOGLEVEL", logging.INFO), + format="%(levelname)s: [%(name)s] -> %(message)s", +) + +syslog = logging.handlers.SysLogHandler() + +logger = logging.getLogger(__name__) +logger.addHandler(syslog) +logger.debug("Init at debug") + + +def start_daemon( + base_uri: str, + name: str, + timeout: int = None, + proxy: T.Optional[T.Tuple[str, int]] = None, + headless: bool = True, + debug: bool = False, +) -> Operator: + """ + Spawns the Operator object. + """ + optional_args = {} # type: T.Dict[str, T.Any] + if timeout is not None: + optional_args["timeout"] = timeout + if proxy is not None: + optional_args["proxy"] = proxy + optional_args["headless"] = headless + optional_args["debug"] = debug + logger.info("Started Operator.") + logger.debug( + "Operator parameters: " + "{base_uri: %s, name: %s, timeout: %s, proxy: %s, headless: %s, debug: %s}", + base_uri, + name, + timeout, + proxy, + headless, + debug, + ) + return Operator(base_uri=base_uri, name=name, **optional_args) + + +def drop_operator(op: Operator) -> None: + """ + Stops the Operator. + """ + logger.debug("Stopping Operator...") + logger.debug("Operator id: %s", id(op)) + op.quit() + logger.info("Stopped Operator.") + + +def reload_operator(op: Operator) -> None: + """ + Reload the Operator using the same object. + """ + op_focus = op.switch_to.active_element + op_focus.send_keys(Keys.CONTROL + "Escape") + op.delete_all_cookies() + logger.info("Operator session cleaned.") + + +def operator_status(op: Operator, debug: bool = False) -> T.Dict[str, str]: + """ + Get the current operator status. + """ + status = {} # type: T.Dict[str, str] + status["is_logged_in"] = str(op.logged_in) + status["is_checked_in"] = str(op.checked_in) + if debug: + status["interal_state"] = str(dir(op)) + logger.info("Status: %s", status) + + return status + + +def daemon_process( + fifo_path: str, + working_dir: str, + umask: int, + pidfile: lockfile.FileLock, + base_uri: str, + name: str, + timeout: int = None, + proxy: T.Optional[T.Tuple[str, int]] = None, + headless: bool = True, + debug: bool = False, + foreground: bool = False, +) -> None: + """ + The daemon function. + """ + op = start_daemon(base_uri, name, timeout, proxy, headless, debug) + logger.debug("Started operator id: %s", id(op)) + context = daemon.DaemonContext( + working_directory=working_dir, umask=umask, pidfile=pidfile + ) + _stop = functools.partial(drop_operator, op) + _reload = functools.partial(reload_operator, op) + _status = functools.partial(operator_status, op) + context.signal_map = { + signal.SIGTERM: _stop, + signal.SIGHUP: "terminate", + signal.SIGUSR1: _reload, + signal.SIGUSR2: _status, + } + cmd_map = {"stop": _stop, "reload": _reload, "status": _status} + + if foreground: + listen_commands(fifo_path, cmd_map) + return + with context: + listen_commands(fifo_path, cmd_map) + + +class Fifo(object): + """ + Iterator to continuously read the command fifo. + """ + + def __init__(self, fifopath: str, mode: str = "r") -> None: + try: + os.mkfifo(fifopath) + logger.info("Control pipe opened at: %s", fifopath) + except OSError as err: + if err.errno != errno.EEXIST: + logger.critical("Could not open control pipe at: %s", fifopath) + raise + self.fh = open(fifopath, mode) + + def __iter__(self): # pragma: noqa + return self + + def __next__(self) -> str: + data = self.fh.read() + # This is a never ending iterator. We should exit the + # loop independently. + return data + + def __del__(self) -> None: + self.fh.close() + + def __enter__(self) -> T.TextIO: + return self.fh + + def __exit__(self, *args) -> None: + self.fh.close() + + def write(self, data) -> None: + self.fh.write(data) + self.fh.flush() + + +def listen_commands(fifopath: str, cmd_map: T.Dict[str, T.Callable]) -> None: + """ + The main loop to listen incoming commands. + """ + while True: + fifo = Fifo(fifopath, "r") + for command in fifo: + if command == "stop": + logger.debug("Received command: %s", "stop") + cmd_map["stop"]() + return + elif command == "reload": + logger.debug("Received command: %s", "reload") + cmd_map["reload"]() + elif command.startswith("status"): + logger.debug("Received command: %s", "status") + status = cmd_map["status"]() + logger.debug("Status: %s", status) + jenc = json.JSONEncoder() + with Fifo(command.split("|")[-1], "w") as rfifo: + rfifo.write(jenc.encode(status)) diff --git a/setup.py b/setup.py index 3155724..3777e44 100644 --- a/setup.py +++ b/setup.py @@ -30,7 +30,12 @@ BIN_PATH = "bin/geckodriver" with open("README.md") as readme_file: readme = readme_file.read() -requirements = ["Click>=6.0", "selenium>=3.141.0"] +requirements = [ + "Click>=7.0", + "selenium>=3.141.0", + "lockfile>=0.12.2", + "python-daemon>=2.2.3", +] setup_requirements = [] # type: T.List[str] @@ -242,7 +247,7 @@ setup( "install": CustomInstallCommand, "bdist_wheel": CustomBDistWheel, }, - entry_points={"console_scripts": ["bot_z=bot_z.cli:main"]}, + entry_points={"console_scripts": ["bot_z=bot_z.cli:cli"]}, package_data={"bot_z": ["bin/geckodriver"]}, include_package_data=True, install_requires=requirements,