From 4a05b05ace9493a67c2d9f65648653a9766865da Mon Sep 17 00:00:00 2001
From: Davte <davte@libero.it>
Date: Sun, 12 Apr 2020 00:25:04 +0200
Subject: [PATCH] Show progress bar in client

---
 src/client.py | 41 ++++++++++++++++++++++++++++++++++++++++-
 src/server.py |  6 ++++++
 2 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/src/client.py b/src/client.py
index 24a668c..80112a0 100644
--- a/src/client.py
+++ b/src/client.py
@@ -6,6 +6,7 @@ import os
 import random
 import ssl
 import string
+import sys
 
 
 class Client:
@@ -102,6 +103,8 @@ class Client:
         writer.write(
             f"s|{self.token}|{file_name}|{file_size}\n".encode('utf-8')
         )
+        self.set_file_information(file_name=file_name,
+                                  file_size=file_size)
         await writer.drain()
         # Wait for server start signal
         while 1:
@@ -156,6 +159,7 @@ class Client:
             while not os.path.isfile(file_path):
                 await asyncio.sleep(.5)
         logging.info("Sending file...")
+        bytes_sent = 0
         with open(file_path, 'rb') as file_to_send:
             while not self.stopping:
                 output_data = file_to_send.read(self.buffer_chunk_size)
@@ -172,6 +176,23 @@ class Client:
                     logging.info('Server closed the connection.')
                     self.stop()
                     break
+                bytes_sent += self.buffer_chunk_size
+                new_progress = min(
+                    int(bytes_sent / self.file_size * 100),
+                    100
+                )
+                progress_showed = (new_progress // 10) * 10
+                sys.stdout.write(
+                    f"\t\t\tSending `{self.file_name}`: "
+                    f"{'#' * (progress_showed // 10)}"
+                    f"{'.' * ((100 - progress_showed) // 10)}\t"
+                    f"{new_progress}% completed "
+                    f"({min(bytes_sent, self.file_size) // 1000} "
+                    f"of {self.file_size // 1000} KB)\r"
+                )
+                sys.stdout.flush()
+        sys.stdout.write('\n')
+        sys.stdout.flush()
         writer.close()
         return
 
@@ -218,11 +239,29 @@ class Client:
             file_path += '.enc'
         logging.info("Receiving file...")
         with open(file_path, 'wb') as file_to_receive:
+            bytes_received = 0
             while not self.stopping:
                 input_data = await reader.read(self.buffer_chunk_size)
+                bytes_received += self.buffer_chunk_size
+                new_progress = min(
+                    int(bytes_received / self.file_size * 100),
+                    100
+                )
+                progress_showed = (new_progress // 10) * 10
+                sys.stdout.write(
+                    f"\t\t\tReceiving `{self.file_name}`: "
+                    f"{'#' * (progress_showed // 10)}"
+                    f"{'.' * ((100 - progress_showed) // 10)}\t"
+                    f"{new_progress}% completed "
+                    f"({min(bytes_received, self.file_size) // 1000} "
+                    f"of {self.file_size // 1000} KB)\r"
+                )
+                sys.stdout.flush()
                 if not input_data:
                     break
                 file_to_receive.write(input_data)
+        sys.stdout.write('\n')
+        sys.stdout.flush()
         logging.info("File received.")
         if self.password:
             logging.info("Decrypting file...")
@@ -257,7 +296,7 @@ class Client:
         if file_name is not None:
             self._file_name = file_name
         if file_size is not None:
-            self._file_size = file_size
+            self._file_size = int(file_size)
 
 
 def get_action(action):
diff --git a/src/server.py b/src/server.py
index 9f82669..7c2dc6b 100644
--- a/src/server.py
+++ b/src/server.py
@@ -71,6 +71,9 @@ class Server:
                 if connection_token not in self.buffers:
                     break
                 self.buffers[connection_token].append(input_data)
+            except ConnectionResetError as e:
+                logging.error(e)
+                break
             except Exception as e:
                 logging.error(e, exc_info=True)
 
@@ -96,6 +99,9 @@ class Server:
                     break
                 writer.write(input_data)
                 await writer.drain()
+            except ConnectionResetError as e:
+                logging.error(e)
+                break
             except Exception as e:
                 logging.error(e, exc_info=True)
                 errors += 1