#!/usr/bin/env python3 import argparse, os, sys directory = os.path.dirname(os.path.realpath(__file__)) home = os.path.expanduser("~") default_data = os.path.join(home, ".bitbot") arg_parser = argparse.ArgumentParser( description="BitBot CLI control utility") arg_parser.add_argument("--socket", "-s", help="Location of bitbotd's unix domain socket", default=os.path.join(default_data, "bot.sock")) arg_parser.add_argument("command") args, unknown = arg_parser.parse_known_args() def _die(code: int, s: str): sys.stderr.write("%s\n" % s) sys.exit(code) SIMPLE = ["rehash", "reload", "stop"] if args.command == "log": arg_parser.add_argument("--level", "-l", help="Log level", default="INFO") elif args.command == "command": arg_parser.add_argument("subcommand") elif args.command in SIMPLE: pass else: _die(1, "Unknown command '%s'" % args.command) args = arg_parser.parse_args() import json, socket, signal def _sigint(_1, _2): print("") sys.exit(0) signal.signal(signal.SIGINT, _sigint) sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) try: sock.connect(args.socket) except ConnectionRefusedError: _die(2, "Failed to connect to BitBot instance. Is bitbotd running?") def _send(s): sock.send(("%s\n" % s).encode("utf8")) # protocol: [arg [arg ..]] _send("0 version 0") if args.command == "log": _send("1 log %s" % args.level) elif args.command == "command": _send("1 command %s" % args.subcommand) elif args.command in SIMPLE: _send("1 %s" % args.command) read_buffer = b"" while True: data = sock.recv(1024) if not data: break data = read_buffer+data lines = [line.strip(b"\r") for line in data.split(b"\n")] read_buffer = lines.pop(-1) for line in lines: line = json.loads(line) if line["action"] == "log": print(line["data"]) elif line["action"] == "ack" and line["data"]: print(line["data"])