colony/colony/cli.py

import argparse
import json
import sys
import time
from datetime import datetime

from colony.core.storages.chat import ChatStorage
from colony.settings import setting
from colony.transport import control

CLEAR = "\033[H\033[2J"


class CliError(Exception):
    pass


def call(args, action, **params):
    try:
        reply = control.request({"action": action, "args": params}, path=args.socket, timeout=args.timeout)
    except OSError as e:
        raise CliError(f"cannot reach the node at {args.socket}: {e}")
    if reply.get("status") != "ok":
        raise CliError(f"the node answered with an error: {reply.get('error')}")
    return reply["result"]


def table(rows):
    widths = [max(len(str(row[i])) for row in rows) for i in range(len(rows[0]))]
    return "\n".join("  ".join(str(cell).ljust(width) for cell, width in zip(row, widths)).rstrip() for row in rows)


def usage(stats, name):
    used, total = stats.get(f"{name}_used_gb"), stats.get(f"{name}_total_gb")
    return f"{used:.1f}/{total:.1f} GB" if used is not None and total is not None else "-"


def render_status(cluster):
    members = cluster["members"]
    nodes = [["NODE", "ADDRESS", "STATE", "CPU", "RAM", "DISK", "ROLES"]]
    for member in members:
        stats = member.get("stats") or {}
        nodes.append([
            member["hostname"] + (" (this node)" if member["self"] else ""),
            f"{member['ip']}:{member['port']}",
            ("alive" if member["alive"] else "DOWN") + (" via gateway" if member.get("via") else ""),
            f"{stats['cpu_percent']:.0f}%" if "cpu_percent" in stats else "-",
            usage(stats, "mem"),
            usage(stats, "disk"),
            " ".join(member["roles"]) or "-",
        ])

    roles = [["ROLE", "PRIORITY", "HOLDERS"]]
    for role in cluster["roles"]:
        holders = [member["hostname"] for member in members if member["alive"] and role["name"] in member["roles"]]
        empty = "off (on demand)" if role.get("on_demand") else "UNMET"
        roles.append([role["name"], role["priority"], ", ".join(holders) or empty])

    return table(nodes) + "\n\n" + table(roles)


def render_chat(messages, peer):
    if not messages:
        return f"No messages with {peer} yet."
    return "\n".join(f"{datetime.fromisoformat(m['at']).astimezone():%H:%M}  {'you' if m['direction'] == 'out' else peer}: {m['text']}" for m in messages)


def find_node(cluster, name):
    matches = [
        member for member in cluster["members"]
        if name.lower() in (member["hostname"].lower(), f"{member['ip']}:{member['port']}", f"{member['hostname']}@{member['ip']}".lower())
    ]
    if len(matches) != 1:
        raise CliError(f"{'no node' if not matches else 'more than one node'} matches {name!r}")
    return matches[0]


def show(args, render):
    if args.json or args.raw or not sys.stdout.isatty():
        print(render())
        return
    try:
        while True:
            try:
                text = render()
            except CliError as e:
                text = f"colonyctl: {e}\nretrying..."
            print(CLEAR + text + "\n\nCtrl-C to quit.", flush=True)
            time.sleep(args.interval)
    except KeyboardInterrupt:
        print()


def status(args):
    def render():
        cluster = call(args, "status")
        return json.dumps(cluster, indent=2) if args.json else render_status(cluster)
    show(args, render)


def add_role(args, name):
    cluster = call(args, "status")
    on_demand = {role["name"].lower(): role["name"] for role in cluster["roles"] if role.get("on_demand")}
    if name.lower() not in on_demand:
        raise CliError(f"{name!r} is not an on-demand role, pick one of: {', '.join(on_demand.values())}")
    role = on_demand[name.lower()]

    reply = call(args, "add_role", role=role)
    if not reply["accepted"]:
        raise CliError(f"the node did not take {role}: {reply['reason']}")
    print(f"{role} is running ({reply['reason']})")
    if role == "WebView":
        me = next(member for member in cluster["members"] if member["self"])
        print(f"Web view: http://{me['ip']}:{setting('web.port')}")


def chat(args):
    peer = find_node(call(args, "status"), args.node)
    peer_id = f"{peer['hostname']}@{peer['ip']}"

    if args.text is not None:
        call(args, "send_chat", peer=peer_id, text=args.text)
        print(f"Sent to {peer['hostname']}")
    else:
        history = lambda: ChatStorage().messages(peer_id)
        show(args, lambda: json.dumps(history(), indent=2) if args.json else render_chat(history(), peer["hostname"]))


def build_parser():
    common = argparse.ArgumentParser(add_help=False)
    common.add_argument("--socket", default=str(control.CONTROL_SOCKET), help="control socket of the local node")
    common.add_argument("--timeout", type=float, default=5)

    live = argparse.ArgumentParser(add_help=False)
    live.add_argument("--interval", type=float, default=1, help="seconds between refreshes")
    live.add_argument("--raw", action="store_true", help="print once instead of refreshing")
    live.add_argument("--json", action="store_true", help="print the raw data as JSON once")

    parser = argparse.ArgumentParser(prog="colonyctl", description="Inspect and control the local colony node.")
    commands = parser.add_subparsers(dest="command", required=True)

    commands.add_parser("status", parents=[common, live], help="show nodes and roles")

    role = commands.add_parser("role", help="manage roles").add_subparsers(dest="role_command", required=True)
    role.add_parser("add", parents=[common], help="run an on-demand role on this node").add_argument("name")

    commands.add_parser("webview", parents=[common], help="same as: role add WebView")

    chat = commands.add_parser("chat", parents=[common, live], help="show the chat with a node, or send it a message")
    chat.add_argument("node", help="hostname, ip:port or hostname@ip")
    chat.add_argument("text", nargs="?", help="message to send")
    return parser


def main(argv=None):
    args = build_parser().parse_args(argv)
    try:
        if args.command == "status":
            status(args)
        elif args.command == "role":
            add_role(args, args.name)
        elif args.command == "webview":
            add_role(args, "WebView")
        elif args.command == "chat":
            chat(args)
    except CliError as e:
        print(f"colonyctl: {e}", file=sys.stderr)
        return 1
    return 0