Queer European MD passionate about IT
Browse Source

Use original file name

Davte 5 năm trước cách đây
mục cha
commit
8bd0ac76f2
1 tập tin đã thay đổi với 33 bổ sung9 xóa
  1. 33 9
      src/client.py

+ 33 - 9
src/client.py

@@ -201,15 +201,22 @@ class Client:
                                           file_size=file_size)
             elif server_hello == 'start!':
                 break
-            logging.info(f"Server said: {server_hello}")
+            else:
+                logging.info(f"Server said: {server_hello}")
         await self.receive(reader=reader)
 
     async def receive(self, reader: asyncio.StreamReader):
         self._working = True
-        file_path = self.file_path
-        logging.info("Receiving file...")
+        file_path = os.path.join(
+            os.path.abspath(
+                self.file_path
+            ),
+            self.file_name
+        )
+        original_file_path = file_path
         if self.password:
             file_path += '.enc'
+        logging.info("Receiving file...")
         with open(file_path, 'wb') as file_to_receive:
             while not self.stopping:
                 input_data = await reader.read(self.buffer_chunk_size)
@@ -224,7 +231,7 @@ class Client:
                 _subprocess = await asyncio.create_subprocess_shell(
                     "openssl enc -aes-256-cbc "
                     "-md sha512 -pbkdf2 -iter 100000 -salt -d "
-                    f"-in \"{file_path}\" -out \"{self.file_path}\" "
+                    f"-in \"{file_path}\" -out \"{original_file_path}\" "
                     f"-pass pass:{self.password}"
                 )
                 stdout, stderr = await _subprocess.communicate()
@@ -316,7 +323,7 @@ def main():
     cli_parser.add_argument('--path', type=str,
                             default=None,
                             required=False,
-                            help='File path')
+                            help='File path to send / folder path to receive')
     cli_parser.add_argument('--password', '--p', '--pass', type=str,
                             default=None,
                             required=False,
@@ -391,11 +398,28 @@ def main():
         action = get_action(
             input("Do you want to (R)eceive or (S)end a file?\t\t")
         )
+    if (
+            (action == 'send'
+             and not os.path.isfile(os.path.abspath(file_path)))
+            or (action == 'receive'
+                and not os.path.isdir(os.path.abspath(file_path)))
+    ):
+        file_path = None
     while file_path is None:
-        file_path = get_file_path(
-            path=input(f"Enter file to {action}:\t\t\t\t\t\t"),
-            action=action
-        )
+        if action == 'send':
+            file_path = get_file_path(
+                path=input(f"Enter file to send:\t\t\t\t\t\t"),
+                action=action
+            )
+            if not os.path.isfile(os.path.abspath(file_path)):
+                file_path = None
+        elif action == 'receive':
+            file_path = get_file_path(
+                path=input(f"Enter destination folder:\t\t\t\t\t\t"),
+                action=action
+            )
+            if not os.path.isdir(os.path.abspath(file_path)):
+                file_path = None
     if password is None:
         logging.warning(
             "You have provided no password for file encryption.\n"