From 2737e06262d31947fa8b693576a2dd86aa760408 Mon Sep 17 00:00:00 2001
From: minhtrannhat <minh@minhtrannhat.com>
Date: Mon, 27 Nov 2023 15:28:54 -0500
Subject: [PATCH] feat: Get list of files available on server

---
 src/myftp/client.py | 12 +++++++++++-
 src/myftp/server.py | 28 +++++++++++++++++++++++++---
 2 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/src/myftp/client.py b/src/myftp/client.py
index 77ea02e..b6799c7 100644
--- a/src/myftp/client.py
+++ b/src/myftp/client.py
@@ -2,6 +2,7 @@ from socket import socket, AF_INET, SOCK_DGRAM
 from typing import Tuple
 from argparse import ArgumentParser
 import os
+import pickle
 
 # custome type to represent the hostname(server name) and the server port
 Address = Tuple[str, int]
@@ -34,10 +35,11 @@ class UDPClient:
                     "summary",
                     "change",
                     "help",
+                    "list",
                     "bye",
                 ]:
                     print(
-                        f"myftp> - {self.mode} - : Invalid command. Supported commands are put, get, summary, change and help"
+                        f"myftp> - {self.mode} - : Invalid command. Supported commands are put, get, summary, change, list and help"
                     )
 
                 # handling the "bye" command
@@ -46,6 +48,14 @@ class UDPClient:
                     print(f"myftp> - {self.mode} - Session is terminated")
                     break
 
+                elif command == "list":
+                    client_socket.send(command.encode())
+                    encoded_message, server_address = client_socket.recvfrom(4096)
+                    file_list = pickle.loads(encoded_message)
+                    print(f"Received file list from {server_address}: {file_list}")
+                    client_socket.close()
+                    continue
+
                 client_socket.send(command.encode())
                 modified_message = client_socket.recv(2048)
                 print(modified_message.decode())
diff --git a/src/myftp/server.py b/src/myftp/server.py
index ca4ed84..37b583d 100644
--- a/src/myftp/server.py
+++ b/src/myftp/server.py
@@ -1,13 +1,17 @@
 from socket import socket, AF_INET, SOCK_DGRAM
 from argparse import ArgumentParser
 import os
+import pickle
 
 
 class UDPServer:
-    def __init__(self, server_name: str, server_port: int, debug: bool) -> None:
+    def __init__(
+        self, server_name: str, server_port: int, directory_path: str, debug: bool
+    ) -> None:
         self.server_name = server_name
         self.server_port = server_port
         self.mode: str = "UDP"
+        self.directory_path = directory_path
         self.debug = debug
 
     def run(self):
@@ -31,6 +35,14 @@ class UDPServer:
 
                 if message_in_utf8 == "ping":
                     response_message = "pong"
+
+                elif message_in_utf8 == "list":
+                    encoded_message = pickle.dumps(
+                        get_files_in_directory(self.directory_path)
+                    )
+                    self.server_socket.sendto(encoded_message, clientAddress)
+                    continue
+
                 else:
                     response_message = message_in_utf8.upper()
 
@@ -49,7 +61,15 @@ class UDPServer:
             print(f"myftp> - {self.mode} - Closed the server socket\n")
 
 
-def check_directory(path):
+def get_files_in_directory(directory_path: str) -> list[str]:
+    file_list = []
+    for root, _, files in os.walk(directory_path):
+        for file in files:
+            file_list.append(file)
+    return file_list
+
+
+def check_directory(path: str) -> bool:
     if os.path.exists(path):
         if os.path.isdir(path):
             if os.access(path, os.R_OK) and os.access(path, os.W_OK):
@@ -109,7 +129,9 @@ def init():
 
     # UDP client selected here
     if protocol_selection == "2":
-        udp_server = UDPServer(args.ip_addr, args.port_number, args.debug)
+        udp_server = UDPServer(
+            args.ip_addr, args.port_number, args.directory, args.debug
+        )
 
         udp_server.run()