Adding chia peer command

chia/cmds/chia.py

@@ -8,6 +8,7 @@
 from chia.cmds.keys import keys_cmd
 from chia.cmds.netspace import netspace_cmd
 from chia.cmds.passphrase import passphrase_cmd
+from chia.cmds.peer import peer_cmd
 from chia.cmds.plots import plots_cmd
 from chia.cmds.rpc import rpc_cmd
 from chia.cmds.show import show_cmd
@@ -132,6 +133,7 @@ def run_daemon_cmd(ctx: click.Context, wait_for_unlock: bool) -> None:
 cli.add_command(farm_cmd)
 cli.add_command(plotters_cmd)
 cli.add_command(db_cmd)
+cli.add_command(peer_cmd)
 cli.add_command(passphrase_cmd)
 
 

chia/cmds/cmds_util.py

@@ -1,7 +1,23 @@
+from pathlib import Path
+from typing import Any, Callable, Dict, Optional, Type
+
+from chia.rpc.farmer_rpc_client import FarmerRpcClient
+from chia.rpc.full_node_rpc_client import FullNodeRpcClient
+from chia.rpc.harvester_rpc_client import HarvesterRpcClient
+from chia.rpc.rpc_client import RpcClient
+from chia.rpc.wallet_rpc_client import WalletRpcClient
 from chia.types.blockchain_format.sized_bytes import bytes32
 from chia.types.mempool_submission_status import MempoolSubmissionStatus
+from chia.util.default_root import DEFAULT_ROOT_PATH
 from chia.wallet.transaction_record import TransactionRecord
 
+NODE_TYPES: Dict[str, Type[RpcClient]] = {
+    "farmer": FarmerRpcClient,
+    "wallet": WalletRpcClient,
+    "full_node": FullNodeRpcClient,
+    "harvester": HarvesterRpcClient,
+}
+
 
 def transaction_submitted_msg(tx: TransactionRecord) -> str:
     sent_to = [MempoolSubmissionStatus(s[0], s[1], s[2]).to_json_dict_convenience() for s in tx.sent_to]
@@ -10,3 +26,48 @@ def transaction_submitted_msg(tx: TransactionRecord) -> str:
 
 def transaction_status_msg(fingerprint: int, tx_id: bytes32) -> str:
     return f"Run 'chia wallet get_transaction -f {fingerprint} -tx 0x{tx_id}' to get status"
+
+
+async def check_client_connection(rpc_client: RpcClient, node_type: str, rpc_port: int) -> bool:
+    from aiohttp import ClientConnectorError
+
+    try:
+        await rpc_client.healthz()
+    except ClientConnectorError:
+        print(f"Connection error. Check if {node_type.replace('_', ' ')} rpc is running at {rpc_port}")
+        print(f"This is normal if {node_type.replace('_', ' ')} is still starting up")
+        return False
+    return True
+
+
+async def execute_with_any_node(
+    node_type: str,
+    rpc_port: Optional[int],
+    function: Callable,
+    root_path: Path = DEFAULT_ROOT_PATH,
+    *args,
+) -> Any:
+    from chia.util.config import load_config
+    from chia.util.ints import uint16
+
+    if node_type not in NODE_TYPES.keys():
+        print(f"Invalid node type: {node_type}")
+        return
+    # load variables from config file
+    config = load_config(root_path, "config.yaml")
+    self_hostname = config["self_hostname"]
+    if rpc_port is None:
+        rpc_port = config[node_type]["rpc_port"]
+    # select node client type based on string
+    node_client_type = NODE_TYPES[node_type]
+    result = None
+    try:
+        node_client = await node_client_type.create(self_hostname, uint16(rpc_port), root_path, config)
+        # check if we can connect to node, this makes the code cleaner
+        if await check_client_connection(node_client, node_type, rpc_port):
+            result = await function(node_client, config, *args)
+    finally:
+        if node_client is not None:
+            node_client.close()
+            await node_client.await_closed()
+    return result

chia/cmds/peer.py

@@ -0,0 +1,49 @@
+from typing import Optional
+
+import click
+
+from chia.cmds.cmds_util import NODE_TYPES, execute_with_any_node
+from chia.cmds.peer_funcs import peer_async
+
+
+@click.command("peer", short_help="Show, or modify peering connections", no_args_is_help=True)
+@click.option(
+    "-p",
+    "--rpc-port",
+    help=(
+        "Set the port where the farmer, wallet, full node or harvester "
+        "is hosting the RPC interface. See the rpc_port in config.yaml"
+    ),
+    type=int,
+    default=None,
+)
+@click.option(
+    "-c", "--connections", help="List nodes connected to this Full Node", is_flag=True, type=bool, default=False
+)
+@click.option("-a", "--add-connection", help="Connect to another Full Node by ip:port", type=str, default="")
+@click.option(
+    "-r", "--remove-connection", help="Remove a Node by the first 8 characters of NodeID", type=str, default=""
+)
+@click.argument("node_type", type=click.Choice(list(NODE_TYPES.keys())), nargs=1, required=True)
+@click.pass_context
+def peer_cmd(
+    ctx: click.Context,
+    rpc_port: Optional[int],
+    connections: bool,
+    add_connection: str,
+    remove_connection: str,
+    node_type: str,
+) -> None:
+    import asyncio
+
+    asyncio.run(
+        execute_with_any_node(
+            node_type,
+            rpc_port,
+            peer_async,
+            ctx.obj["root_path"],
+            connections,
+            add_connection,
+            remove_connection,
+        )
+    )

chia/cmds/peer_funcs.py

@@ -0,0 +1,117 @@
+from typing import Any, Dict
+
+from chia.rpc.rpc_client import RpcClient
+
+
+async def add_node_connection(rpc_client: RpcClient, add_connection: str) -> None:
+    if ":" not in add_connection:
+        print("Enter a valid IP and port in the following format: 10.5.4.3:8000")
+    else:
+        ip, port = (
+            ":".join(add_connection.split(":")[:-1]),
+            add_connection.split(":")[-1],
+        )
+        print(f"Connecting to {ip}, {port}")
+        try:
+            await rpc_client.open_connection(ip, int(port))
+        except Exception:
+            print(f"Failed to connect to {ip}:{port}")
+
+
+async def remove_node_connection(rpc_client: RpcClient, remove_connection: str) -> None:
+    from chia.server.outbound_message import NodeType
+
+    result_txt = ""
+    if len(remove_connection) != 8:
+        result_txt = "Invalid NodeID. Do not include '.'"
+    else:
+        connections = await rpc_client.get_connections()
+        for con in connections:
+            if remove_connection == con["node_id"].hex()[:8]:
+                print("Attempting to disconnect", "NodeID", remove_connection)
+                try:
+                    await rpc_client.close_connection(con["node_id"])
+                except Exception:
+                    result_txt = f"Failed to disconnect NodeID {remove_connection}"
+                else:
+                    result_txt = (
+                        f"NodeID {remove_connection}... {NodeType(con['type']).name} {con['peer_host']} disconnected"
+                    )
+            elif result_txt == "":
+                result_txt = f"NodeID {remove_connection}... not found"
+    print(result_txt)
+
+
+async def print_connections(rpc_client: RpcClient, trusted_peers: Dict[str, Any]) -> None:
+    import time
+
+    from chia.server.outbound_message import NodeType
+    from chia.util.network import is_trusted_inner
+
+    connections = await rpc_client.get_connections()
+    print("Connections:")
+    print("Type      IP                                     Ports       NodeID      Last Connect" + "      MiB Up|Dwn")
+    for con in connections:
+        last_connect_tuple = time.struct_time(time.localtime(con["last_message_time"]))
+        last_connect = time.strftime("%b %d %T", last_connect_tuple)
+        mb_down = con["bytes_read"] / (1024 * 1024)
+        mb_up = con["bytes_written"] / (1024 * 1024)
+
+        host = con["peer_host"]
+        # Strip IPv6 brackets
+        host = host.strip("[]")
+
+        trusted: bool = is_trusted_inner(host, con["node_id"], trusted_peers, False)
+        # Nodetype length is 9 because INTRODUCER will be deprecated
+        if NodeType(con["type"]) is NodeType.FULL_NODE:
+            peak_height = con.get("peak_height", None)
+            connection_peak_hash = con.get("peak_hash", None)
+            if connection_peak_hash is None:
+                connection_peak_hash = "No Info"
+            else:
+                if connection_peak_hash.startswith(("0x", "0X")):
+                    connection_peak_hash = connection_peak_hash[2:]
+                connection_peak_hash = f"{connection_peak_hash[:8]}..."
+            con_str = (
+                f"{NodeType(con['type']).name:9} {host:38} "
+                f"{con['peer_port']:5}/{con['peer_server_port']:<5}"
+                f" {con['node_id'].hex()[:8]}... "
+                f"{last_connect}  "
+                f"{mb_up:7.1f}|{mb_down:<7.1f}"
+                f"\n                                                 "
+            )
+            if peak_height is not None:
+                con_str += f"-Height: {peak_height:8.0f}    -Hash: {connection_peak_hash}"
+            else:
+                con_str += f"-Height: No Info    -Hash: {connection_peak_hash}"
+            # Only show when Trusted is True
+            if trusted:
+                con_str += f"    -Trusted: {trusted}"
+        else:
+            con_str = (
+                f"{NodeType(con['type']).name:9} {host:38} "
+                f"{con['peer_port']:5}/{con['peer_server_port']:<5}"
+                f" {con['node_id'].hex()[:8]}... "
+                f"{last_connect}  "
+                f"{mb_up:7.1f}|{mb_down:<7.1f}"
+            )
+        print(con_str)
+
+
+async def peer_async(
+    rpc_client: RpcClient,
+    config: Dict[str, Any],
+    show_connections: bool,
+    add_connection: str,
+    remove_connection: str,
+    # trusted_peers: Dict[str, Any],
+) -> None:
+    # Check or edit node connections
+    if show_connections:
+        trusted_peers: Dict[str, Any] = config["full_node"].get("trusted_peers", {})
+        await print_connections(rpc_client, trusted_peers)
+        # if called together with state, leave a blank line
+    if add_connection:
+        await add_node_connection(rpc_client, add_connection)
+    if remove_connection:
+        await remove_node_connection(rpc_client, remove_connection)