diff --git a/kitty/file_transmission.py b/kitty/file_transmission.py index 0d406502a..1b7a9bf3a 100644 --- a/kitty/file_transmission.py +++ b/kitty/file_transmission.py @@ -2,22 +2,48 @@ # vim:fileencoding=utf-8 # License: GPLv3 Copyright: 2021, Kovid Goyal -from enum import Enum, auto from base64 import standard_b64decode +from enum import Enum, auto +from typing import Optional + +from .utils import log_error class Action(Enum): send = auto() data = auto() + end_data = auto() receive = auto() invalid = auto() +class Container(Enum): + zip = auto() + tar = auto() + tgz = auto() + tbz2 = auto() + txz = auto() + none = auto() + + +class Compression(Enum): + zlib = auto() + none = auto() + + +class Encoding(Enum): + base64 = auto() + + class FileTransmissionCommand: action = Action.invalid + container_fmt = Container.none + compression = Compression.none + encoding = Encoding.base64 id: str = '' secret: str = '' + mime: str = '' payload = b'' @@ -35,12 +61,43 @@ def parse_command(data: str) -> FileTransmissionCommand: k, v = x.partition('=')[::2] if k == 'action': ans.action = Action[v] - elif k == 'id': - ans.id = v - elif k == 'secret': - ans.secret = v + elif k == 'container_fmt': + ans.container_fmt = Container[v] + elif k == 'compression': + ans.compression = Compression[v] + elif k == 'encoding': + ans.encoding = Encoding[v] + elif k in ('secret', 'mime', 'id'): + setattr(ans, k, v) if ans.action is Action.invalid: raise ValueError('No valid action specified in file transmission command') return ans + + +class FileTransmission: + + active_cmd: Optional[FileTransmissionCommand] = None + + def __init__(self, window_id: int): + self.window_id = window_id + + def handle_serialized_command(self, data: str) -> None: + try: + cmd = parse_command(data) + except Exception as e: + log_error(f'Failed to parse file transmission command with error: {e}') + return + if self.active_cmd is not None: + if cmd.action not in (Action.data, Action.end_data): + log_error('File transmission command received while another is in flight, aborting') + self.abort_in_flight() + if cmd.action is Action.send: + self.start_send(cmd) + + def start_send(self, cmd: FileTransmissionCommand) -> None: + self.active_cmd = cmd + + def abort_in_flight(self) -> None: + self.active_cmd = None diff --git a/kitty/window.py b/kitty/window.py index de8ef02ff..63ab3fbf8 100644 --- a/kitty/window.py +++ b/kitty/window.py @@ -13,7 +13,7 @@ from gettext import gettext as _ from itertools import chain from typing import ( Any, Callable, Deque, Dict, Iterable, List, NamedTuple, Optional, Pattern, - Sequence, Tuple, Union + Sequence, Tuple, Union, TYPE_CHECKING ) from .child import ProcessDesc @@ -49,6 +49,10 @@ from .utils import ( MatchPatternType = Union[Pattern[str], Tuple[Pattern[str], Optional[Pattern[str]]]] +if TYPE_CHECKING: + from .file_transmission import FileTransmission + + class WindowDict(TypedDict): id: int is_focused: bool @@ -358,6 +362,14 @@ class Window: else: setup_colors(self.screen, opts) + @property + def file_transmission_control(self) -> 'FileTransmission': + ans: Optional['FileTransmission'] = getattr(self, '_file_transmission', None) + if ans is None: + from .file_transmission import FileTransmission + ans = self._file_transmission = FileTransmission(self.id) + return ans + def on_dpi_change(self, font_sz: float) -> None: self.update_effective_padding() @@ -800,8 +812,7 @@ class Window: self.screen.send_escape_code_to_child(DCS, '@kitty-cmd' + json.dumps(response)) def file_transmission(self, data: str) -> None: - from .file_transmission import parse_command - parse_command(data) + self.file_transmission_control.handle_serialized_command(data) def clipboard_control(self, data: str, is_partial: bool = False) -> None: where, text = data.partition(';')[::2]