#!/usr/bin/env python3
import argparse
import json
import random
import socket
import sys

from collections import defaultdict

HOST = 'localhost'
PORT = None
ID_NEXT = random.randint(0, 262144)
JSON = False
EXITSTATUS = 0

class ErrorResponse(RuntimeError):
    pass

def send_request(method, params=None):
    verb = False#not args.q
    with socket.create_connection((HOST, PORT), timeout=10.0) as sock:
        def sndrecv(_id, method, params=None):
            outj = { "id" : _id, "jsonrpc" : "2.0", "method" : method, "params": params or [] }
            msg = json.dumps(outj, indent=None).encode("utf8") + b'\n'
            if verb: print(f"Srv --> {msg[:2048]}")
            sock.send(msg)
            resp = bytearray()
            while b'\n' not in resp:
                resp += sock.recv(4096)
            if verb: print(f"Srv <-- {resp[:2048]}")
            j = json.loads(resp.decode("utf8").strip())
            if j.get("error") or j.get("id") != _id:
                raise ErrorResponse("Error response from Fulcrum:\n\n" + (json.dumps(j.get("error"), indent=4) or j))
            return j.get("result")
        global ID_NEXT
        reqid = ID_NEXT; ID_NEXT += 1
        return sndrecv(reqid, method, params)


def main():
    global HOST, PORT, JSON, EXITSTATUS
    parser = argparse.ArgumentParser(prog="FulcrumAdmin.py", description="Fulcrum CLI admin tool")
    parser.add_argument('-p', type=int, metavar="port", nargs=1, required=True, help=f"Specify the port for the Fulcrum admin RPC service. This is a required argument.")
    parser.add_argument('-j', action='store_true', dest="json", help=f"Print the response from the server as JSON, not as formatted text")
    parser.add_argument('-H', type=str, metavar="host", nargs='?', default=HOST, help=f"Specify the host for the Fulcrum admin RPC service. Defaults to {HOST}.")
    subparsers = parser.add_subparsers(title="command", description="Select from one of the following commands:", dest="command")
    addpeer = subparsers.add_parser('addpeer', help="Add a peer to the server's list of peers")
    addpeer.add_argument('hostname', metavar='hostname', nargs=1, help="Hostname of peer.")
    addpeer.add_argument('-s', metavar='ssl_port', type=int, nargs='?', help="Peer's SSL port.")
    addpeer.add_argument('-t', metavar='tcp_port', type=int, nargs='?', help="Peer's TCP port.")
    ban = subparsers.add_parser('ban', help="Ban clients by ID and/or IP address")
    ban.add_argument('id_or_ip', metavar='ipaddress_or_id', nargs='+', help="Client ID or IP address to ban.")
    banpeer = subparsers.add_parser('banpeer', help="Ban peers by hostname suffix")
    banpeer.add_argument('hostnames', metavar='hostname', nargs='+', help="A hostname or hostname suffix e.g. somehost.com or *some.host.com.")
    bitcoind_throttle = subparsers.add_parser('bitcoind_throttle', help="Query or set server bitcoind_throttle setting")
    bitcoind_throttle.add_argument('param', metavar='param', nargs='*', help='The new desired setting. Specify 3 arguments to set this properly for: high low decay. Omit arguments to query.')
    clients = subparsers.add_parser('clients', help="Print information on all the currently connected clients", aliases=['sessions'])
    getinfo = subparsers.add_parser('getinfo', help="Get server information")
    kick = subparsers.add_parser('kick', help="Kick clients by ID and/or IP address")
    kick.add_argument('id_or_ip', metavar='ipaddress_or_id', nargs='+', help="Client ID or IP addresses to kick.")
    listbanned = subparsers.add_parser('listbanned', help="Print the list of banned IP addresses and peer hostnames", aliases=['banlist'])
    loglevel = subparsers.add_parser('loglevel', help="Set the server's logging verbosity")
    loglevel.add_argument('level', metavar='level', nargs=1, help="One of: 'normal', 'debug', or 'trace'")
    maxbuffer = subparsers.add_parser('maxbuffer', help="Query or set server max_buffer setting")
    maxbuffer.add_argument('bytes', metavar='bytes', type=int, nargs='?', help='The new desired max_buffer setting in bytes. Must be >= 64KB and <= 100MB. If omitted, then this script will just query the current value.')
    peers = subparsers.add_parser('peers', help="Print peering information")
    rmpeer = subparsers.add_parser('rmpeer', help="Remove peers by hostname suffix")
    rmpeer.add_argument('hostnames', metavar='hostname', nargs='+', help="A hostname or hostname suffix e.g. somehost.com or *some.host.com.")
    simdjson = subparsers.add_parser('simdjson', help="Get or set the server's 'simdjson' (JSON parser) setting")
    simdjson.add_argument('enabled', type=int, nargs='?',
                          help='Flag used to enable or disable the simdjson JSON parser on the server (1=enabled,'
                               ' 0=disabled). If this option is omitted, then the current setting is queried.')
    stop = subparsers.add_parser('stop', help="Gracefully shut down the server", aliases=['shutdown'])
    unban = subparsers.add_parser('unban', help="Unban IP addresses")
    unban.add_argument('ips', metavar='ipaddress', nargs='+', help="Specify an existing banned IP address to unban.")
    unbanpeer = subparsers.add_parser('unbanpeer', help="Unban peers by hostname suffix")
    unbanpeer.add_argument('hostnames', metavar='hostname', nargs='+', help="Specify an existing peer ban to unban.")

    args = parser.parse_args()

    HOST = args.H
    PORT, = args.p
    JSON = args.json

    if PORT > 65535:
        sys.exit("Port argument must be < 65536")

    command = args.command
    if command == 'sessions': command = 'clients'  # Is there a better way to do this by referring back to the subparser above?? TODO
    if command == 'shutdown': command = 'stop'
    if command == 'banlist' : command = 'listbanned'
    command_params = tuple()
    response_handler = lambda r: json.dumps(r, indent = 4)  # default handler just pretty-prints the JSON

    if command is None:
        print("Please specify a command to run.\n")
        parser.print_help()
        sys.exit(1)

    elif command == 'getinfo':
        orig_handler = response_handler
        def handler(r):
            # mogrify the uptime field to be more useful to humans
            if r.get('uptime'):
                r['uptime'] = formatTimeField(r['uptime'])
            return orig_handler(r)
        response_handler = handler

    elif command == 'stop' and not JSON:
        def handler(r):
            if isinstance(r, bool) and r:
                return "Fulcrum server is shutting down"
            else:
                global EXITSTATUS
                EXITSTATUS = 1
                return f"Unexpected response from Fulcrum server: {r}"
        response_handler = handler

    elif command == 'clients' and not JSON:
        response_handler = clients_handler

    elif command == 'maxbuffer':
        command_params = [args.bytes] if args.bytes else command_params
        if not JSON:
            def handler(x):
                if command_params:
                    return f"Server max_buffer setting -> {x}"
                else:
                    return f"Server max_buffer setting is: {x}"
            response_handler = handler

    elif command == 'simdjson':
        command_params = [bool(args.enabled)] if args.enabled is not None else command_params
        if not JSON:
            def handler(x):
                if command_params:
                    # set
                    assert isinstance(x, bool)
                    if x:
                        verb = "enabled" if command_params[0] else "disabled"
                        return f"Server simdjson parser {verb} successfully"
                    else:
                        verb = "enable" if command_params[0] else "disable"
                        return f"Failed to {verb} server simdjson parser"
                else:
                    # query
                    if x is not None:
                        return json.dumps(x, indent=4)
                    else:
                        return "Server simdjson parser is disabled"
            response_handler = handler

    elif command == 'bitcoind_throttle':
        command_params = args.param if len(args.param) else command_params
        if not JSON:
            def handler(x):
                names = defaultdict(lambda: "???")
                names.update({ 0: "high", 1: "low", 2: "decay" })
                if command_params:
                    return f"Server bitcoind_throttle setting -> " + ', '.join(names[n] + ' = ' + str(i) for n,i in enumerate(x))
                else:
                    return "Server bitcoind_throttle setting is: " + ', '.join(names[n] + ' = ' + str(i) for n,i in enumerate(x))
            response_handler = handler

    elif command in ('kick', 'ban', 'banpeer', 'unban', 'unbanpeer', 'rmpeer', 'loglevel'):
        extratxt = ''
        if command in ('kick', 'ban'):
            command_params = args.id_or_ip
        elif command in ('banpeer', 'unbanpeer', 'rmpeer'):
            command_params = args.hostnames
            if command == 'rmpeer':
                extratxt = ('''

Note: Removal of 'Good' peers is not guaranteed to keep them from re-peering
with the server in the near future. If you want the specified peer(s) to never
possibly peer with this Fulcrum server, then please use the 'banpeer' command.
''')
        elif command == 'unban':
            command_params = args.ips
        elif command == 'loglevel':
            levels = { 'normal': 0, 'debug' : 1, 'trace' : 2}
            if args.level[0] not in levels:
                print("level argument must be one of:", ', '.join(levels.keys()))
                sys.exit(1)
            command_params = [levels[args.level[0]]]
            extratxt = f' -> {args.level[0]}'

        if not JSON:
            def handler(r):
                if isinstance(r, bool) and r:
                    return f"{command} command submitted for: " + ', '.join([str(x) for x in command_params]) + extratxt
                else:
                    global EXITSTATUS
                    EXITSTATUS = 1
                    return f"Unexpected response from Fulcrum server: {r}"
            response_handler = handler

    elif command == 'listbanned' and not JSON:
        response_handler = listbanned_handler

    elif command == 'peers' and not JSON:
        response_handler = peers_handler

    elif command == 'addpeer':
        if not args.s and not args.t:
            print("addpeer requires at least one of the two port arguments (-s, -t), or both.")
            sys.exit(1)
        host = args.hostname[0]
        command_params = {
            'host' : host,
            'ssl'  : args.s or 0,
            'tcp'  : args.t or 0,
        }
        if not JSON:
            def handler(r):
                if isinstance(r, bool) and r:
                    return f"{command} command submitted for: {host}"
                else:
                    global EXITSTATUS
                    EXITSTATUS = 1
                    return f"Unexpected response from Fulcrum server: {r}"
            response_handler = handler

    try:
        print(response_handler(send_request(command, command_params)))
        sys.exit(EXITSTATUS)
    except OSError as e:
        print(f"Error communicating with the admin RPC port at {HOST}:{PORT}\n\n    {e}\n")
        sys.exit(1)
    except ErrorResponse as e:
        print(f"{e}")
        sys.exit(1)



def clients_handler(r):
    def badResp():
        global EXITSTATUS
        EXITSTATUS = 1
        return f"Unexpected resposne from Fulcrum server: {r!r}"
    if not isinstance(r, (list, tuple)):
        return badResp()
    lines = []
    line = ("ID","IP:PORT","Typ","UAgent","ProtocolVer","Subs","HdrSub?","ReqRcv","RespSent","RecvBytes","SentBytes","TxsSent","Notifs","ErrorCt","Elapsed")
    maxfields = defaultdict(int)
    for i,c in enumerate(line):
        maxfields[i] = max(maxfields[i], len(c))
    lines.append(line)
    for serverdict in r:
        if not isinstance(serverdict, dict):
            return badResp()
        for server_name, subdict in serverdict.items():
            if not isinstance(subdict, dict):
                # in case we add stuff to here that is not a server dict some day
                continue;
            is_wss = server_name.lower().startswith('wsss')
            is_ws = not is_wss and server_name.lower().startswith('wss')
            if not is_wss and not is_ws:
                is_ssl = server_name.lower().startswith('ssl')
                typnam = 'SSL' if is_ssl else 'TCP'
            else:
                typnam = 'WS' if is_ws else 'WSS'

            for client in subdict.get('clients', []):
                for cname, cdict in client.items():
                    line = (
                        cdict.get('id', -1),
                        cdict.get('remote', '?'),
                        typnam,
                        cdict.get('userAgent', 'Unk'),
                        cdict.get('version', ['?'])[-1],
                        cdict.get('nSubscriptions', -1),
                        'Y' if cdict.get('isSubscribedToHeaders') else 'N',
                        cdict.get('nRequestsRcv', -1),
                        cdict.get('nResultsSent', -1),
                        cdict.get('nBytesReceived', -1),
                        cdict.get('nBytesSent', -1),
                        cdict.get('nTxSent', -1),
                        cdict.get('nNotificationsSent', -1),
                        cdict.get('nErrorsSent', -1),
                        formatTimeField( cdict.get('connectedTime', '-') )
                    )
                    line = [str(x) for x in line]
                    for i,c in enumerate(line):
                        maxfields[i] = max(maxfields[i], len(c))
                    lines.append(line)
    for i,line in enumerate(list(lines)):
        line = list(line)
        for j,c in enumerate(line):
            line[j] = c.ljust(maxfields[j])
        lines[i] = '  '.join(line)
    return '\n'.join(lines) + '\n'


def listbanned_handler(r):
    def badResp():
        global EXITSTATUS
        EXITSTATUS = 1
        return f"Unexpected resposne from Fulcrum server: {r!r}"
    if not isinstance(r, (dict,)):
        return badResp()
    lines = []
    line = ("IP","AgeSecs","RejectedConnections",)
    maxfields = defaultdict(int)
    for i,c in enumerate(line):
        maxfields[i] = max(maxfields[i], len(c))
    lines.append(line)
    for ipaddr, subdict in r.get('Banned_IPAddrs', {}).items():
        if not isinstance(subdict, dict):
            return badResp()
        line = (
            ipaddr,
            subdict.get('age_secs', -1),
            subdict.get('connections_rejected', 0),
        )
        line = [str(x) for x in line]
        for i,c in enumerate(line):
            maxfields[i] = max(maxfields[i], len(c))
        lines.append(line)
    for i,line in enumerate(list(lines)):
        line = list(line)
        for j,c in enumerate(line):
            line[j] = c.ljust(maxfields[j])
        lines[i] = '  '.join(line)
    clientPart = '~Client Bans~\n' + '\n'.join(lines)
    maxFields = defaultdict(int)
    lines = []
    line = ("HostName","AgeSecs",)
    for i,c in enumerate(line):
        maxfields[i] = max(maxfields[i], len(c))
    lines.append(line)
    for hostname, subdict in r.get('Banned_Peers', {}).items():
        if not isinstance(subdict, dict):
            return badResp()
        line = (
            hostname,
            subdict.get('age_secs', -1),
        )
        line = [str(x) for x in line]
        for i,c in enumerate(line):
            maxfields[i] = max(maxfields[i], len(c))
        lines.append(line)
    for i,line in enumerate(list(lines)):
        line = list(line)
        for j,c in enumerate(line):
            line[j] = c.ljust(maxfields[j])
        lines[i] = '  '.join(line)
    peerPart = '~Peer Bans~\n' + '\n'.join(lines)
    return clientPart + '\n\n' + peerPart


def peers_handler(r):
    def badResp():
        global EXITSTATUS
        EXITSTATUS = 1
        return f"Unexpected resposne from Fulcrum server: {r!r}"

    if not isinstance(r, (dict,)):
        return badResp()
    lines = []
    line = ("Hostname","IP","Status","TCP","SSL","Version","ProtoMin","ProtoMax","Elapsed","Message","RetryPeriod")
    maxfields = defaultdict(int)
    for i,c in enumerate(line):
        maxfields[i] = max(maxfields[i], len(c))
    lines.append(line)
    retryTimes = {
        "bad" : r.get('activeTimers', {}).get("badPeerRetry", 0),
        "failed" : r.get('activeTimers', {}).get("failedPeerRetry", 0),
    }
    # peers (good, connected peers)
    peerdict = r.get('peers')
    if not isinstance(peerdict, dict):
        return badResp()
    for hostname, d in peerdict.items():
        line = (
            hostname,
            d.get('addr','-'),
            'Good' if d.get('verified') == True else 'Verifying',
            d.get('tcp_port') if d.get('tcp_port') else '-',
            d.get('ssl_port') if d.get('ssl_port') else '-',
            d.get('server_version', 'Unk'),
            d.get('protocol_min', '-'),
            d.get('protocol_max', '-'),
            formatTimeField( d.get('connectedTime', '-') ),
            '-', '-'
        )
        line = [str(x) for x in line]
        for i,c in enumerate(line):
            maxfields[i] = max(maxfields[i], len(c))
        lines.append(line)
    # next, do the failed, bad, and queued peers
    for status in ('failed', 'bad', 'queued'):
        d = r.get(status)
        if not d:
            continue  # skip empties
        if not isinstance(d, dict):
            return badResp()
        for hostname, l in d.items():
            try:
                ip, tcp, ssl, subver, pver, time, msg = l
            except ValueError:
                pass
            line = (
                hostname, ip, status.title(), tcp if tcp else '-', ssl if ssl else '-',
                subver, pver, pver, formatTimeField(time), msg, '~' + str(retryTimes.get(status, 0)//1000) + 's'
            )
            line = [str(x) for x in line]
            for i,c in enumerate(line):
                maxfields[i] = max(maxfields[i], len(c))
            lines.append(line)

    for i,line in enumerate(list(lines)):
        line = list(line)
        for j,c in enumerate(line):
            line[j] = c.ljust(maxfields[j])
        lines[i] = '  '.join(line)
    return '\n'.join(lines) + '\n'


def formatTimeField(s):
    if s is None:
        return '-'  # transform None to '-' since it takes up fewer characters
    secs_per_hour = 60.0 * 60.0
    if isinstance(s, str):
        to_secs_factor = 1.0
        if s.endswith("hours"):  # one of the "xx hours" fields
            to_secs_factor = secs_per_hour
        elif s.endswith("secs"):
            pass
        else:
            return s # unknown field type, just return it verbatim
        try:
            s = float(s.split()[0].strip()) * to_secs_factor  # transform hours to seconds
        except (ValueError, TypeError, IndexError):
            return s
        del to_secs_factor
    # at this point s should be in seconds.. if not, give up and just return it
    if not isinstance(s, (int, float)):
        return s
    # now turn seconds to minutes, hours, days, months, years
    secs_per_day = secs_per_hour * 24.0
    secs_per_year = 365.0 * secs_per_day
    secs_per_month = 30.0 * secs_per_day
    secs_per_minute = 60.0
    if s >= secs_per_year:
        return f'{s/secs_per_year:0.3f} years'
    elif s >= secs_per_month:
        return f'{s/secs_per_month:0.3f} months'
    elif s >= secs_per_day:
        return f'{s/secs_per_day:0.2f} days'
    elif s >= secs_per_hour:
        return f'{s/secs_per_hour:0.2f} hours'
    elif s >= secs_per_minute:
        return f'{s/secs_per_minute:0.1f} mins'
    else:
        return f'{s:0.1f} secs'


if __name__ == '__main__':
    main()
