8000 Support for asyncssh server integration. · mxr/python-prompt-toolkit@954e166 · GitHub
[go: up one dir, main page]
Skip to content

Commit 954e166

Browse files
jonathanslendersjonathanslenders
committed
Support for asyncssh server integration.
1 parent f38ad93 commit 954e166

File tree

3 files changed

+230
-0
lines changed

3 files changed

+230
-0
lines changed

‎examples/ssh/asyncssh-server.py

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
#!/usr/bin/env python
2+
"""
3+
Example of running a prompt_toolkit application in an asyncssh server.
4+
"""
5+
import asyncio
6+
import logging
7+
8+
import asyncssh
9+
10+
from prompt_toolkit.shortcuts.dialogs import yes_no_dialog, input_dialog
11+
from prompt_toolkit.shortcuts.prompt import PromptSession
12+
from prompt_toolkit.shortcuts import print_formatted_text
13+
from prompt_toolkit.shortcuts import ProgressBar
14+
from prompt_toolkit.contrib.ssh import PromptToolkitSSHServer
15+
16+
from pygments.lexers.html import HtmlLexer
17+
18+
from prompt_toolkit.lexers import PygmentsLexer
19+
20+
from prompt_toolkit.completion import WordCompleter
21+
22+
animal_completer = WordCompleter([
23+
'alligator', 'ant', 'ape', 'bat', 'bear', 'beaver', 'bee', 'bison',
24+
'butterfly', 'cat', 'chicken', 'crocodile', 'dinosaur', 'dog', 'dolphin',
25+
'dove', 'duck', 'eagle', 'elephant', 'fish', 'goat', 'gorilla', 'kangaroo',
26+
'leopard', 'lion', 'mouse', 'rabbit', 'rat', 'snake', 'spider', 'turkey',
27+
'turtle',
28+
], ignore_case=True)
29+
30+
31+
async def interact() -> None:
32+
"""
33+
The application interaction.
34+
35+
This will run automatically in a prompt_toolkit AppSession, which means
36+
that any prompt_toolkit application (dialogs, prompts, etc...) will use the
37+
SSH channel for input and output.
38+
"""
39+
prompt_session = PromptSession()
40+
41+
# Alias 'print_formatted_text', so that 'print' calls go to the SSH client.
42+
print = print_formatted_text
43+
44+
print('We will be running a few prompt_toolkit applications through this ')
45+
print('SSH connection.\n')
46+
47+
# Simple progress bar.
48+
with ProgressBar() as pb:
49+
for i in pb(range(50)):
50+
await asyncio.sleep(.1)
51+
52+
# Normal prompt.
53+
text = await prompt_session.prompt_async("(normal prompt) Type something: ")
54+
print("You typed", text)
55+
56+
# Prompt with auto completion.
57+
text = await prompt_session.prompt_async(
58+
"(autocompletion) Type an animal: ", completer=animal_completer)
59+
print("You typed", text)
60+
61+
# prompt with syntax highlighting.
62+
text = await prompt_session.prompt_async("(HTML syntax highlighting) Type something: ",
63+
lexer=PygmentsLexer(HtmlLexer))
64+
print("You typed", text)
65+
66+
# Show yes/no dialog.
67+
await prompt_session.prompt_async('Showing yes/no dialog... [ENTER]')
68+
await yes_no_dialog("Yes/no dialog", "Running over asyncssh").run_async()
69+
70+
# Show input dialog
71+
await prompt_session.prompt_async('Showing input dialog... [ENTER]')
72+
await input_dialog("Input dialog", "Running over asyncssh").run_async()
73+
74+
75+
def main(port=8222):
76+
# Set up logging.
77+
logging.basicConfig()
78+
logging.getLogger().setLevel(logging.DEBUG)
79+
80+
loop = asyncio.get_event_loop()
81+
loop.run_until_complete(
82+
asyncssh.create_server(
83+
lambda: PromptToolkitSSHServer(interact),
84+
"",
85+
port,
86+
server_host_keys=["/etc/ssh/ssh_host_e 6D40 cdsa_key"],
87+
)
88+
)
89+
loop.run_forever()
90+
91+
92+
if __name__ == "__main__":
93+
main()
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
from .server import PromptToolkitSession, PromptToolkitSSHServer
2+
3+
__all__ = [
4+
'PromptToolkitSession',
5+
'PromptToolkitSSHServer',
6+
]
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,131 @@
1+
"""
2+
Utility for running a prompt_toolkit application in an asyncssh server.
3+
"""
4+
import asyncio
5+
import traceback
6+
from typing import Awaitable, Callable, Optional, TextIO, cast
7+
8+
import asyncssh
9+
from prompt_toolkit.application.current import AppSession, create_app_session
10+
from prompt_toolkit.data_structures import Size
11+
from prompt_toolkit.input.posix_pipe import PosixPipeInput
12+
from prompt_toolkit.output.vt100 import Vt100_Output
13+
14+
__all__ = [
15+
'PromptToolkitSession',
16+
'PromptToolkitSSHServer',
17+
]
18+
19+
20+
class PromptToolkitSession(asyncssh.SSHServerSession):
21+
def __init__(self, interact: Callable[[], Awaitable[None]]) -> None:
22+
self.interact = interact
23+
self._chan = None
24+
self.app_session: Optional[AppSession] = None
25+
26+ F438
# PipInput object, for sending input in the CLI.
27+
# (This is something that we can use in the prompt_toolkit event loop,
28+
# but still write date in manually.)
29+
self._input = PosixPipeInput()
30+
31+
# Output object. Don't render to the real stdout, but write everything
32+
# in the SSH channel.
33+
class Stdout:
34+
def write(s, data):
35+
if self._chan is not None:
36+
self._chan.write(data.replace('\n', '\r\n'))
37+
38+
def flush(s):
39+
pass
40+
41+
self._output = Vt100_Output(cast(TextIO, Stdout()),
42+
self._get_size, write_binary=False)
43+
44+
def _get_size(self) -> Size:
45+
"""
46+
Callable that returns the current `Size`, required by Vt100_Output.
47+
"""
48+
if self._chan is None:
49+
return Size(rows=20, columns=79)
50+
else:
51+
width, height, pixwidth, pixheight = self._chan.get_terminal_size()
52+
return Size(rows=height, columns=width)
53+
54+
def connection_made(self, chan):
55+
self._chan = chan
56+
57+
def shell_requested(self) -> bool:
58+
return True
59+
60+
def session_started(self) -> None:
61+
asyncio.create_task(self._interact())
62+
63+
async def _interact(self) -> None:
64+
if self._chan is None:
65+
# Should not happen.
66+
raise Exception('`_interact` called before `connection_made`.')
67+
68+
# Disable the line editing provided by asyncssh. Prompt_toolkit
69+
# provides the line editing.
70+
self._chan.set_line_mode(False)
71+
72+
with create_app_session(input=self._input, output=self._output) as session:
73+
self.app_session = session
74+
try:
75+
await self.interact()
76+
except BaseException:
77+
traceback.print_exc()
78+
finally:
79+
# Close the connection.
80+
self._chan.close()
81+
82+
def terminal_size_changed(self, width, height, pixwidth, pixheight):
83+
# Send resize event to the current application.
84+
if self.app_session and self.app_session.app:
85+
self.app_session.app._on_resize()
86+
87+
def data_received(self, data, datatype):
88+
self._input.send_text(data)
89+
90+
91+
class PromptToolkitSSHServer(asyncssh.SSHServer):
92+
"""
93+
Run a prompt_toolkit application over an asyncssh server.
94+
95+
This takes one argument, an `interact` function, which is called for each
96+
connection. This should be an asynchronous function that runs the
97+
prompt_toolkit applications. This function runs in an `AppSession`, which
98+
means that we can have multiple UI interactions concurrently.
99+
100+
Example usage:
101+
102+
.. code:: python
103+
104+
async def interact() -> None:
105+
await yes_no_dialog("my title", "my text").run_async()
106+
107+
prompt_session = PromptSession()
108+
text = await prompt_session.prompt_async("Type something: ")
109+
print_formatted_text('You said: ', text)
110+
111+
server = PromptToolkitSSHServer(interact=interact)
112+
loop = get_event_loop()
113+
loop.run_until_complete(
114+
asyncssh.create_server(
115+
lambda: MySSHServer(interact),
116+
"",
117+
port,
118+
server_host_keys=["/etc/ssh/..."],
119+
)
120+
)
121+
loop.run_forever()
122+
"""
123+
def __init__(self, interact: Callable[[], Awaitable[None]]) -> None:
124+
self.interact = interact
125+
126+
def begin_auth(self, username):
127+
# No authentication.
128+
return False
129+
130+
def session_requested(self) -> PromptToolkitSession:
131+
return PromptToolkitSession(self.interact)