I'm trying to do file transferring via sockets between two machines. I don't have any problems with transferring but I want to be able to resume inputting another command in my server after the transfer process and print that download successful after finished. This is my Server code:
import shlex
import socket
server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
server.bind(("10.87.13.44", 9999))
server.listen()
print("[+] Waiting for incoming connections")
server, addr = server.accept()
print("[+] Got a connection from " + str(addr))
def system_commands_results():
com_result = server.recv(1024).decode()
data_size = int(com_result)
com_result = b""
while len(com_result) < data_size:
data = server.recv(4096)
if not data:
break
com_result += data
print(com_result.decode())
def receive_and_write_file(file_path):
with open(file_path, "wb") as file:
received_file_result = server.recv(1024)
received_size = int(received_file_result.decode())
received_file_result = b""
while len(received_file_result) < received_size:
received_data = server.recv(4096)
if not received_data:
break
file.write(received_data)
return "[+] Download successful"
while True:
try:
while True:
command = input(">> ")
command = shlex.split(command)
if command[0] == "download":
path = command[1]
server.send("download".encode())
server.send(path.encode())
status = receive_and_write_file(path)
print(status)
result = server.recv(1024).decode()
print(result)
elif command[0] == "exit":
server.send("exit".encode())
exit()
elif command[0] == "cd" and len(command) > 1:
server.send("cd ".encode())
change_to = command[1]
server.send(change_to.encode())
result = server.recv(1024).decode()
print(result)
else:
server.send(command[0].encode())
system_commands_results()
except Exception as e:
print(f"[-] Error: {e}")
Client code:
import os
import subprocess
import socket
client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
client.connect(("10.87.13.44", 9999))
def execute_system_command(command):
com_result = subprocess.check_output(command, shell=True, stderr=subprocess.STDOUT, text=True)
result_size = str(len(com_result)).encode('utf-8')
client.send(result_size)
chunk_size = 4096
for i in range(0, len(com_result), chunk_size):
client.send(com_result[i:i + chunk_size].encode('utf-8'))
def read_and_send_file(file_path):
with open(file_path, "rb") as file:
file_content = file.read()
file_size = str(len(file_content)).encode()
client.send(file_size)
chunk = 4096
for j in range(0, len(file_content), chunk):
client.send(file_content[j:j + chunk])
return "Finished"
def change_working_directory(dir_path):
os.chdir(dir_path)
return f"[+] Changing working directory to {os.getcwd()}"
while True:
try:
received_command = client.recv(1024).decode()
if received_command == "download":
path = client.recv(1024).decode()
result = read_and_send_file(path)
client.send(result.encode())
elif received_command == "exit":
exit()
elif received_command == "cd ":
change_to = client.recv(1024).decode()
result = change_working_directory(change_to)
client.send(result.encode())
else:
execute_system_command(received_command)
except Exception as e:
client.send(str(e).encode())
Right now after the transfer is completed my program just waits and I can't continue to input commands
It's, unfortunately, not 100% straightforward as discussed in the other answer's comments; you can't trust .recv() to always return as many bytes as you desire, so you'll need some sort of framing. For instance, HTTP, being a text-based protocol, has framing based on a double \r\n\r\n sequence to determine that this is the end of the client's headers, and what follows is a body (whose length is specified in the Content-length header.
However, this isn't HTTP, so we'll come up with a small TLV (type-length-value) scheme (2 bytes for message type, 4 bytes of message length); you can imagine other message types than just 1 (ping from client to server) and 2 (pong from server to client), such as 3 for "file name" and 4 for "file contents", 5 for "please run this command" and 6 for "here's that command's results" or what-have-you.
This example is self-contained in that it has two threads, a client thread, and server thread, and they talk to each other over a socket.
import socket
import struct
import threading
import time
port = 12340
def read_n_bytes(sock, n: int) -> bytes:
buf = b""
while len(buf) < n:
buf += sock.recv(n - len(buf))
return buf
def server():
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
sock.bind(("localhost", port))
sock.listen(1)
conn, addr = sock.accept()
handle_single_connection(conn, addr)
def read_message_header(conn):
# if we receive less than 6 bytes, struct.unpack will also raise
buf = conn.recv(6)
if not buf:
raise RuntimeError("Connection closed")
msg_type, msg_len = struct.unpack(">HI", buf)
return msg_type, msg_len
def write_message(sock, msg_type: int, content: bytes):
sock.sendall(struct.pack(">HI", msg_type, len(content)))
sock.sendall(content)
def handle_single_connection(conn, addr):
print("Server: Connected by", addr)
while True:
try:
msg_type, msg_len = read_message_header(conn)
except RuntimeError: # connection closed
print("Server: Connection closed")
break
if msg_type == 1: # ping
ping_content = read_n_bytes(conn, msg_len)
print("Server: Received ping:", ping_content)
# Respond with a pong with the ping content reversed, for funsies
write_message(conn, 2, ping_content[::-1])
else:
print("Server: Unknown message type, dying")
break
def client():
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
sock.connect(("localhost", port))
for x in range(1, 10):
message = f"hello! {1 << x}".encode("utf-8")
write_message(sock, 1, message)
msg_type, msg_len = read_message_header(sock)
if msg_type == 2: # pong
print("Client: Received pong:", read_n_bytes(sock, msg_len))
else:
print("Client: Unknown message type, dying")
break
time.sleep(0.5)
print("Client: closing connection")
def main():
threading.Thread(target=server).start()
time.sleep(0.1)
threading.Thread(target=client).start()
if __name__ == "__main__":
main()