mirror of
https://github.com/kovidgoyal/kitty
synced 2026-07-24 09:18:08 +02:00
Work on supporting streaming remote commands with passwords
This commit is contained in:
@@ -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}'
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user