AsyncSSH passphrase callable tester

callable_test.py

#! /usr/bin/env python3

import asyncio
import asyncssh
import sys
from textual import on
from textual.app import App
from textual.screen import ModalScreen
from textual.widgets import Label, Button, Input, Header, Footer, Log

def key_is_encrypted(fname):
	try:
		asyncssh.read_private_key(fname)
		return False
	except asyncssh.KeyImportError:
		return True

async def passphrase_query(app, fname):
	# our callback is currently being called even on unencrypted keys, so
	# check if the key is encrypted and only ask for a passphrase if it is
	if not key_is_encrypted(fname):
		return ''

	# for some reason app.push_screen_wait must be called from a worker
	screen = app.push_screen_wait(PassphraseQueryScreen(fname))
	worker = app.run_worker(screen)
	passphrase = await worker.wait()
	return passphrase

async def ssh_connect_worker(app, btn, log):
	async def passphrase_query_with_app(fname):
		return await passphrase_query(app, fname)

	async with asyncssh.connect(sys.argv[1], passphrase=passphrase_query_with_app) as conn:
		result = await conn.run('ls /')
		log.clear()
		log.write_line(result.stdout)
	btn.disabled = False

# this is to show that main thread isn't blocked
async def flip_label_worker(lbl):
	cnt = 0
	while True:
		lbl.update(f'{cnt:.1f}')
		await asyncio.sleep(0.2)
		cnt += 0.2

class PassphraseQueryScreen(ModalScreen):
	CSS = '''
PassphraseQueryScreen {
	align: center middle;
}
Input {
	width: 30;
}
'''
	def __init__(self, filename):
		super().__init__()
		self.filename = filename

	def compose(self):
		yield Label(f'Enter passphrase for {self.filename}')
		yield Input(password=True)

	def on_input_submitted(self, event):
		password = event.value
		self.dismiss(password)

class SSHApp(App):
	def compose(self):
		yield Header()
		yield Label('[loading]', id='status-lbl')
		yield Button('connect', id='connect')
		yield Log(id='ssh-log')
		yield Footer()

	@on(Button.Pressed, '#connect')
	async def connect(self):
		btn = self.query_one('#connect')
		btn.disabled = True
		log = self.query_one('#ssh-log')
		self.run_worker(ssh_connect_worker(self, btn, log))

	async def on_mount(self):
		lbl = self.query_one('#status-lbl')
		self.run_worker(flip_label_worker(lbl))

app = SSHApp()
app.run()