Update servefiles.

3ds/servefiles.py

@@ -1,4 +1,6 @@
 #!/usr/bin/env python
+# coding: utf-8 -*-
+
 import os
 import socket
 import struct
@@ -18,69 +20,73 @@ except ImportError:
     from urllib.parse import urljoin, quote
     from urllib.request import pathname2url
 
-if len(sys.argv) < 3:
-	print("Usage: " + sys.argv[0] + " <ip> <file/directory> [host ip]")
+if len(sys.argv) < 3 or len(sys.argv) > 5:
+    print('Usage: ' + sys.argv[0] + ' <target ip> <file / directory> [host ip] [host port]')
     sys.exit(1)
 
-ip = sys.argv[1]
-directory = sys.argv[2]
+accepted_extension = ('.cia', '.tik', '.cetk')
+target_ip = sys.argv[1]
+target_path = sys.argv[2]
+hostPort = 8080 # Default value
 
-if not os.path.exists(directory):
-	print(directory + ": No such file or directory.")
+if not os.path.exists(target_path):
+    print(target_path + ': No such file or directory.')
     sys.exit(1)
 
 if len(sys.argv) >= 4:
     hostIp = sys.argv[3]
+    if len(sys.argv) == 5:
+        hostPort = int(sys.argv[4])
 else:
+    print('Detecting host IP...')
     hostIp = [(s.connect(('8.8.8.8', 53)), s.getsockname()[0], s.close()) for s in [socket.socket(socket.AF_INET, socket.SOCK_DGRAM)]][0][1]
 
-print("Preparing data...")
+print('Preparing data...')
+baseUrl = hostIp + ':' + str(hostPort) + '/'
 
-baseUrl = hostIp + ":8080/"
-payload = ""
-
-if os.path.isfile(directory):
-	payload += baseUrl + quote(os.path.basename(directory))
-	directory = os.path.dirname(directory)
-else:
-	for file in [ file for file in next(os.walk(directory))[2] if file.endswith(('.cia', '.tik')) ]:
-		payload += baseUrl + quote(file) + "\n"
-
-if len(payload) == 0:
-	print("No files to serve.")
+if os.path.isfile(target_path):
+    if target_path.endswith(accepted_extension):
+        file_list_payload = baseUrl + quote(os.path.basename(target_path))
+        directory = os.path.dirname(target_path)  # get file directory
+    else:
+        print('Unsupported file extension. Supported extensions are: ' + accepted_extension)
         sys.exit(1)
 
-payloadBytes = payload.encode("ascii")
+else:
+    directory = target_path  # it's a directory
+    file_list_payload = ''  # init the payload before adding lines
+    for file in [file for file in next(os.walk(target_path))[2] if file.endswith(accepted_extension)]:
+        file_list_payload += baseUrl + quote(file) + '\n'
 
-if not directory == "":
-	os.chdir(directory)
+if len(file_list_payload) == 0:
+    print('No files to serve.')
+    sys.exit(1)
 
-print("")
-print("URLS:")
-print(payload)
-print("")
+file_list_payloadBytes = file_list_payload.encode('ascii')
 
-print("Opening HTTP server on port 8080...")
+if directory and directory != '.':  # doesn't need to move if it's already the current working directory
+    os.chdir(directory)  # set working directory to the right folder to be able to serve files
 
-server = TCPServer(("", 8080), SimpleHTTPRequestHandler)
+print('\nURLs:')
+print(file_list_payload + '\n')
+
+print('Opening HTTP server on port ' + str(hostPort))
+server = TCPServer(('', hostPort), SimpleHTTPRequestHandler)
 thread = threading.Thread(target=server.serve_forever)
 thread.start()
 
 try:
-	print("Sending URL(s) to " + ip + ":5000...")
-
+    print('Sending URL(s) to ' + target_ip + ' on port 5000...')
     sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-	sock.connect((ip, 5000))
-	sock.sendall(struct.pack('!L', len(payloadBytes)) + payloadBytes)
+    sock.connect((target_ip, 5000))
+    sock.sendall(struct.pack('!L', len(file_list_payloadBytes)) + file_list_payloadBytes)
     while len(sock.recv(1)) < 1:
         time.sleep(0.05)
-
     sock.close()
 except Exception as e:
-	print("Error: " + str(e))
+    print('An error occurred: ' + str(e))
     server.shutdown()
     sys.exit(1)
 
-print("Shutting down HTTP server...")
-
+print('Shutting down HTTP server...')
 server.shutdown()