Work on supporting streaming remote commands with passwords

This commit is contained in:
Kovid Goyal
2022-08-31 20:03:51 +05:30
parent d7985689c9
commit 364533b1ed
14 changed files with 206 additions and 85 deletions

View File

@@ -1,11 +1,12 @@
#!/usr/bin/env python
# License: GPLv3 Copyright: 2020, Kovid Goyal <kovid at kovidgoyal.net>
import tempfile
from contextlib import suppress
from dataclasses import dataclass
from typing import (
TYPE_CHECKING, Any, Callable, Dict, FrozenSet, Iterable, Iterator, List,
NoReturn, Optional, Set, Tuple, Type, Union, cast
TYPE_CHECKING, Any, Callable, Dict, FrozenSet, Iterable, Iterator,
List, NoReturn, Optional, Set, Tuple, Type, Union, cast
)
from kitty.cli import get_defaults_from_seq, parse_args, parse_option_spec
@@ -29,6 +30,16 @@ class NoResponse:
pass
class NamedTemporaryFile:
name: str = ''
def __enter__(self) -> None: ...
def __exit__(self, exc: Any, value: Any, tb: Any) -> None: ...
def close(self) -> None: ...
def write(self, data: bytes) -> None: ...
def flush(self) -> None: ...
class RemoteControlError(Exception):
pass
@@ -51,6 +62,11 @@ class UnknownLayout(ValueError):
hide_traceback = True
class StreamError(ValueError):
hide_traceback = True
class PayloadGetter:
def __init__(self, cmd: 'RemoteCommand', payload: Dict[str, Any]):
@@ -245,6 +261,35 @@ class ArgsHandling:
raise TypeError(f'Unknown args handling for cmd: {cmd_name}')
class StreamInFlight:
def __init__(self) -> None:
self.stream_id = ''
self.tempfile: Optional[NamedTemporaryFile] = None
def handle_data(self, stream_id: str, data: bytes) -> Union[AsyncResponse, NamedTemporaryFile]:
from ..remote_control import close_active_stream
if stream_id != self.stream_id:
close_active_stream(self.stream_id)
if self.tempfile is not None:
self.tempfile.close()
self.tempfile = None
self.stream_id = stream_id
if self.tempfile is None:
t: NamedTemporaryFile = cast(NamedTemporaryFile, tempfile.NamedTemporaryFile(suffix='.png'))
self.tempfile = t
else:
t = self.tempfile
if data:
t.write(data)
return AsyncResponse()
close_active_stream(self.stream_id)
self.stream_id = ''
self.tempfile = None
t.flush()
return t
class RemoteCommand:
Args = ArgsHandling
@@ -253,7 +298,6 @@ class RemoteCommand:
desc: str = ''
args: ArgsHandling = ArgsHandling()
options_spec: Optional[str] = None
no_response: bool = False
response_timeout: float = 10. # seconds
string_return_is_error: bool = False
defaults: Optional[Dict[str, Any]] = None
@@ -262,10 +306,12 @@ class RemoteCommand:
protocol_spec: str = ''
argspec = args_count = args_completion = ArgsHandling()
field_to_option_map: Optional[Dict[str, str]] = None
reads_streaming_data: bool = False
def __init__(self) -> None:
self.desc = self.desc or self.short_desc
self.name = self.__class__.__module__.split('.')[-1].replace('_', '-')
self.stream_in_flight = StreamInFlight()
def fatal(self, msg: str) -> NoReturn:
if running_in_kitty():
@@ -342,6 +388,12 @@ class RemoteCommand:
def cancel_async_request(self, boss: 'Boss', window: Optional['Window'], payload_get: PayloadGetType) -> None:
pass
def handle_streamed_data(self, data: bytes, payload_get: PayloadGetType) -> Union[NamedTemporaryFile, AsyncResponse]:
stream_id = payload_get('stream_id')
if not stream_id or not isinstance(stream_id, str):
raise StreamError('No stream_id in rc payload')
return self.stream_in_flight.handle_data(stream_id, data)
def cli_params_for(command: RemoteCommand) -> Tuple[Callable[[], str], str, str, str]:
return (command.options_spec or '\n').format, command.args.spec, command.desc, f'{appname} @ {command.name}'

View File

@@ -100,8 +100,9 @@ not interpreted for escapes. If stdin is a terminal, you can press :kbd:`Ctrl+D`
Path to a file whose contents you wish to send. Note that in this case the file contents
are sent as is, not interpreted for escapes.
'''
no_response = True
args = RemoteCommand.Args(spec='[TEXT TO SEND]', json_field='data', special_parse='+session_id:parse_send_text(io_data, args)')
is_asynchronous = True
reads_streaming_data = True
def message_to_kitty(self, global_opts: RCOptions, opts: 'CLIOptions', args: ArgsType) -> PayloadType:
limit = 1024

View File

@@ -3,15 +3,14 @@
import imghdr
import os
import tempfile
from base64 import standard_b64decode, standard_b64encode
from typing import IO, TYPE_CHECKING, Dict, Optional
from typing import TYPE_CHECKING, Optional
from kitty.types import AsyncResponse
from .base import (
MATCH_WINDOW_OPTION, ArgsType, Boss, CmdGenerator, PayloadGetType,
PayloadType, RCOptions, RemoteCommand, ResponseType, Window
MATCH_WINDOW_OPTION, ArgsType, Boss, CmdGenerator, NamedTemporaryFile,
PayloadGetType, PayloadType, RCOptions, RemoteCommand, ResponseType, Window
)
if TYPE_CHECKING:
@@ -65,8 +64,7 @@ failed, the command will exit with a success code.
''' + '\n\n' + MATCH_WINDOW_OPTION
args = RemoteCommand.Args(spec='PATH_TO_PNG_IMAGE', count=1, json_field='data', special_parse='!read_window_logo(args[0])', completion={
'files': ('PNG Images', ('*.png',))})
images_in_flight: Dict[str, IO[bytes]] = {}
is_asynchronous = True
reads_streaming_data = True
def message_to_kitty(self, global_opts: RCOptions, opts: 'CLIOptions', args: ArgsType) -> PayloadType:
if len(args) != 1:
@@ -98,34 +96,26 @@ failed, the command will exit with a success code.
def response_from_kitty(self, boss: Boss, window: Optional[Window], payload_get: PayloadGetType) -> ResponseType:
data = payload_get('data')
img_id = payload_get('async_id')
if data != '-':
if img_id not in self.images_in_flight:
self.images_in_flight[img_id] = tempfile.NamedTemporaryFile(suffix='.png')
if data:
self.images_in_flight[img_id].write(standard_b64decode(data))
return AsyncResponse()
windows = self.windows_for_payload(boss, window, payload_get)
os_windows = tuple({w.os_window_id for w in windows if w})
layout = payload_get('layout')
if data == '-':
path = None
tfile = NamedTemporaryFile()
else:
f = self.images_in_flight.pop(img_id)
path = f.name
f.flush()
q = self.handle_streamed_data(standard_b64decode(data) if data else b'', payload_get)
if isinstance(q, AsyncResponse):
return q
path = q.name
tfile = q
try:
boss.set_background_image(path, os_windows, payload_get('configured'), layout)
with tfile:
boss.set_background_image(path, os_windows, payload_get('configured'), layout)
except ValueError as err:
err.hide_traceback = True # type: ignore
raise
return None
def cancel_async_request(self, boss: 'Boss', window: Optional['Window'], payload_get: PayloadGetType) -> None:
async_id = payload_get('async_id')
self.images_in_flight.pop(async_id, None)
set_background_image = SetBackgroundImage()

View File

@@ -4,15 +4,14 @@
import imghdr
import os
import tempfile
from base64 import standard_b64decode, standard_b64encode
from typing import IO, TYPE_CHECKING, Dict, Optional
from typing import TYPE_CHECKING, Optional
from kitty.types import AsyncResponse
from .base import (
MATCH_WINDOW_OPTION, ArgsType, Boss, CmdGenerator, PayloadGetType,
PayloadType, RCOptions, RemoteCommand, ResponseType, Window
MATCH_WINDOW_OPTION, ArgsType, Boss, CmdGenerator, NamedTemporaryFile,
PayloadGetType, PayloadType, RCOptions, RemoteCommand, ResponseType, Window
)
if TYPE_CHECKING:
@@ -61,8 +60,7 @@ failed, the command will exit with a success code.
'''
args = RemoteCommand.Args(spec='PATH_TO_PNG_IMAGE', count=1, json_field='data', special_parse='!read_window_logo(args[0])', completion={
'files': ('PNG Images', ('*.png',))})
images_in_flight: Dict[str, IO[bytes]] = {}
is_asynchronous = True
reads_streaming_data = True
def message_to_kitty(self, global_opts: RCOptions, opts: 'CLIOptions', args: ArgsType) -> PayloadType:
if len(args) != 1:
@@ -94,26 +92,22 @@ failed, the command will exit with a success code.
def response_from_kitty(self, boss: Boss, window: Optional[Window], payload_get: PayloadGetType) -> ResponseType:
data = payload_get('data')
img_id = payload_get('async_id')
if data != '-':
if img_id not in self.images_in_flight:
self.images_in_flight[img_id] = tempfile.NamedTemporaryFile(suffix='.png')
if data:
self.images_in_flight[img_id].write(standard_b64decode(data))
return AsyncResponse()
if data == '-':
path = ''
else:
f = self.images_in_flight.pop(img_id)
path = f.name
f.flush()
alpha = float(payload_get('alpha', '-1'))
position = payload_get('position') or ''
for window in self.windows_for_match_payload(boss, window, payload_get):
if window:
window.set_logo(path, position, alpha)
if data == '-':
path = ''
tfile = NamedTemporaryFile()
else:
q = self.handle_streamed_data(standard_b64decode(data) if data else b'', payload_get)
if isinstance(q, AsyncResponse):
return q
path = q.name
tfile = q
with tfile:
for window in self.windows_for_match_payload(boss, window, payload_get):
if window:
window.set_logo(path, position, alpha)
return None