From b484fcf9ede862914055701ce308cfeadc609997 Mon Sep 17 00:00:00 2001
From: Fabio Manganiello <blacklight86@gmail.com>
Date: Mon, 20 Jan 2020 15:39:23 +0100
Subject: [PATCH] - New implementation of esp.file_get that also works with
 binary files

- New communication protocol for downloaded chunks based on queues
  instead of events.
---
 platypush/plugins/esp/__init__.py          | 30 ++++++++++++++--------
 platypush/plugins/esp/models/connection.py | 23 +++++++++--------
 2 files changed, 33 insertions(+), 20 deletions(-)

diff --git a/platypush/plugins/esp/__init__.py b/platypush/plugins/esp/__init__.py
index fe646fe62c..9fa06c9c67 100644
--- a/platypush/plugins/esp/__init__.py
+++ b/platypush/plugins/esp/__init__.py
@@ -1,4 +1,5 @@
 import base64
+import io
 import os
 import threading
 
@@ -1532,22 +1533,31 @@ print([b for b in os.urandom({size})])
         return self.execute(code, **kwargs).output
 
     @action
-    def file_get(self, file: str, **kwargs) -> str:
+    def file_get(self, file: str, binary: bool = False, timeout: Optional[float] = 60.0, **kwargs) -> str:
         """
         Get the content of a file on the board.
-        NOTE: It only works with non-binary files.
 
         :param file: File name/path to get from the device.
-        :param kwargs: Parameters to pass to :meth:`platypush.plugins.esp.EspPlugin.execute`.
+        :param binary: If True, then the base64-encoded content of the file will be returned.
+        :param timeout: File transfer timeout (default: one minute).
+        :param kwargs: Parameters to pass to :meth:`platypush.plugins.esp.EspPlugin.connect`.
         """
-        file = self.string_quote(file)
-        code = '''
-import os
-with open('{file}', 'r') as f:
-    print(f.read())
-'''.format(file=file)
+        device = self._get_device(**kwargs)
+        host = device['host']
+        port = device['port']
+        self.connect(host=host, port=port, password=device['password'])
+        conn = self._get_connection(host=host, port=port)
 
-        return self.execute(code, **kwargs).output
+        with io.BytesIO() as buffer:
+            conn.file_download(file, buffer, timeout=timeout)
+            data = buffer.getvalue()
+
+        if binary:
+            data = base64.encodebytes(data).decode()
+        else:
+            data = data.decode()
+
+        return data
 
     @action
     def file_upload(self, source: str, destination: Optional[str] = None, timeout: Optional[float] = 60.0, **kwargs):
diff --git a/platypush/plugins/esp/models/connection.py b/platypush/plugins/esp/models/connection.py
index 7fc225c0ad..3b90049266 100644
--- a/platypush/plugins/esp/models/connection.py
+++ b/platypush/plugins/esp/models/connection.py
@@ -1,5 +1,6 @@
 import enum
 import logging
+import queue
 import os
 import re
 import threading
@@ -50,7 +51,7 @@ class Connection:
         self._file_transfer_ack_received = threading.Event()
         self._file_transfer_request_successful = True
         self._file_transfer_successful = True
-        self._downloaded_chunk = None
+        self._downloaded_chunks = queue.Queue()
         self._password_requested = False
         self._running_cmd = None
         self._received_echo = None
@@ -203,7 +204,7 @@ class Connection:
     def on_file_download_start(self):
         self.logger.info('Starting file download')
         self._file_transfer_successful = False
-        self._downloaded_chunk = None
+        self._downloaded_chunks = queue.Queue()
         self.state = self.State.DOWNLOADING_FILE
         self._file_transfer_ack_received.clear()
         self.ws.send(b'\x00', opcode=websocket.ABNF.OPCODE_BINARY)
@@ -218,10 +219,10 @@ class Connection:
         if size == 0:
             # End of file
             self.logger.info('File download completed')
-            self._downloaded_chunk = None
+            self._downloaded_chunks.put(None)
             self.on_file_download_completed()
         else:
-            self._downloaded_chunk = data
+            self._downloaded_chunks.put(data)
             self.ws.send(b'\x00', opcode=websocket.ABNF.OPCODE_BINARY)
 
         self._download_chunk_ready.set()
@@ -237,14 +238,16 @@ class Connection:
         self.state = self.State.SENDING_FILE_TRANSFER_RESPONSE
 
     def get_downloaded_chunks(self, timeout: Optional[float] = None):
-        block_ok = self._download_chunk_ready.wait(timeout=timeout)
+        while True:
+            try:
+                chunk = self._downloaded_chunks.get(timeout=timeout)
+            except queue.Empty:
+                self.on_timeout('File download timed out')
 
-        while self._downloaded_chunk and block_ok:
-            yield self._downloaded_chunk
-            block_ok = self._download_chunk_ready.wait(timeout=timeout)
+            if chunk is None:
+                break
 
-        if not block_ok:
-            self.on_timeout('File download timed out')
+            yield chunk
 
     def wait_ready(self):
         connected = self._connected.wait(timeout=self.connect_timeout)