mmgen-node-tools/mmnode-peerblocks

#!/usr/bin/env python3
#
# mmgen = Multi-Mode GENerator, command-line Bitcoin cold storage solution
# Copyright (C)2013-2020 The MMGen Project <mmgen@tuta.io>
#
# This program is free software: you can redistribute it and/or modify it under
# the terms of the GNU General Public License as published by the Free Software
# Foundation, either version 3 of the License, or (at your option) any later
# version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
# details.
#
# You should have received a copy of the GNU General Public License along with
# this program.  If not, see <http://www.gnu.org/licenses/>.

"""
mmgen-peerblocks: List blocks in flight, disconnect stalling nodes
"""

import asyncio
from mmgen.common import *

opts_data = {
	'text': {
		'desc': 'List blocks in flight, disconnect stalling nodes',
		'usage':   '[opts]',
		'options': """
-h, --help      Print this help message
--, --longhelp  Print help message for long options (common options)
"""
	}
}

opts.init(opts_data)

async def inflight_display(rpc):

	def gen_peers(peerinfo):
		global min_height
		min_height = None
		for d in peerinfo:
			if 'inflight' in d and d['inflight']:
				blks = d['inflight']
				if not min_height or min_height > blks[0]:
					min_height = blks[0]
				blks_trunc = ' '.join(map(str,blks))[:term_width-6].split()
				trim = blks_trunc[-1] != str(blks[len(blks_trunc)-1])
				blks = blks[:len(blks_trunc)-trim]
			else:
				blks = []
			yield { 'id': d['id'], 'data': blks }

	def gen_line(peer):
		for blk in peer['data']:
			if blk == min_height:
				yield RED + str(blk) + RESET
			else:
				yield COLORS[blk % 10] + str(blk) + RESET

	from mmgen.term import get_terminal_size
	term_width = get_terminal_size()[0]

	count = 1
	while True:
		info = await rpc.call('getpeerinfo')

		msg_r(CUR_HOME+ERASE_ALL+CUR_HOME)
		msg(f'ACTIVE PEERS ({len(info)}) - poll {count}')

		for peer in gen_peers(info):
			sys.stderr.write('\r{} {:>3}: {}\n'.format(
				ERASE_LINE,
				peer['id'],
				' '.join(gen_line(peer)) ))

		msg_r(ERASE_ALL+'Hit ENTER for disconnect prompt: ')
		await asyncio.sleep(2)
		count += 1

async def do_inflight(rpc):
	task = asyncio.ensure_future(inflight_display(rpc)) # Python 3.7+: create_task()
	from select import select

	while True:
		key = select([sys.stdin], [], [], 0.1)[0]
		if key:
			sys.stdin.read(1)
			task.cancel()
			break
		await asyncio.sleep(0.1)

	try:
		await task
	except asyncio.CancelledError:
		pass

async def do_disconnect_menu(rpc):

	while True:
		peerinfo = await rpc.call('getpeerinfo')
		ids = [str(d['id']) for d in peerinfo]

		msg_r(CUR_HOME+ERASE_ALL+CUR_HOME)
		msg(f'ACTIVE PEERS ({len(peerinfo)})')

		if peerinfo:
			msg('\n'.join([f"{d['id']:>3}: {d['addr']:30} {d['subver']}" for d in peerinfo]))

		reply = input("Peer number to disconnect, ENTER to quit menu, 'u' to update peer list> ")

		if reply == '':
			return
		elif reply == 'u':
			msg(f'Updating peer list')
			await asyncio.sleep(0.5)
		elif reply in ids:
			addr = peerinfo[ids.index(reply)]['addr']
			msg(f'Disconnecting peer {reply} ({addr})')
			try:
				await rpc.call('disconnectnode',addr)
			except RPCFailure:
				msg(f'Unable to disconnect peer {addr}')
			await asyncio.sleep(1.5)
		else:
			msg(f'{reply!r}: invalid peer number')
			await asyncio.sleep(0.5)

async def main():

	msg_r(CUR_HOME+ERASE_ALL)

	from mmgen.rpc import rpc_init
	rpc = await rpc_init()

	while True:
		await do_inflight(rpc)
		await do_disconnect_menu(rpc)

RED,RESET = ('\033[31m','\033[0m')
COLORS = ['\033[38;5;%s;1m' % c for c in (238,240,242,244,246,247,249,251,253,255)]
ERASE_ALL,ERASE_LINE,CUR_HOME,CUR_HIDE,CUR_SHOW = (
	'\033[J','\033[K','\033[H','\033[?25l','\033[?25h')

try:
	run_session(main(),do_rpc_init=False)
except:
	from subprocess import run
	run(['stty','sane'])
	msg('')