More work on file transfer protocol

kitty/file_transmission.py

@@ -2,22 +2,48 @@
 # vim:fileencoding=utf-8
 # License: GPLv3 Copyright: 2021, Kovid Goyal <kovid at kovidgoyal.net>
 
-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':
+        elif k == 'container_fmt':
-            ans.id = v
+            ans.container_fmt = Container[v]
-        elif k == 'secret':
+        elif k == 'compression':
-            ans.secret = v
+            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

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
+        self.file_transmission_control.handle_serialized_command(data)
-        parse_command(data)
 
     def clipboard_control(self, data: str, is_partial: bool = False) -> None:
         where, text = data.partition(';')[::2]