forked from openwpm/OpenWPM
-
Notifications
You must be signed in to change notification settings - Fork 2
/
SocketInterface.py
179 lines (161 loc) · 6.53 KB
/
SocketInterface.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
import Queue
import threading
import traceback
import socket
import struct
import json
import dill
#TODO - Implement a cleaner shutdown for server socket
# see: https://stackoverflow.com/questions/1148062/python-socket-accept-blocks-prevents-app-from-quitting
class serversocket:
"""
A server socket to recieve and process string messages
from client sockets to a central queue
"""
def __init__(self, verbose=False):
self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self.sock.bind(('localhost', 0))
self.sock.listen(10) # queue a max of n connect requests
self.verbose = verbose
self.queue = Queue.Queue()
if self.verbose:
print "Server bound to: " + str(self.sock.getsockname())
def start_accepting(self):
""" Start the listener thread """
thread = threading.Thread(target=self._accept, args=())
thread.daemon = True # stops from blocking shutdown
thread.start()
def _accept(self):
""" Listen for connections and pass handling to a new thread """
while True:
(client, address) = self.sock.accept()
thread = threading.Thread(target=self._handle_conn, args=(client, address))
thread.daemon = True
thread.start()
def _handle_conn(self, client, address):
"""
Recieve messages and pass to queue. Messages are prefixed with
a 4-byte integer to specify the message length and 1-byte character
to indicate the type of serialization applied to the message.
Supported serialization formats:
'n' : no serialization
'd' : dill pickle
'j' : json
"""
if self.verbose:
print "Thread: " + str(threading.current_thread()) + " connected to: " + str(address)
try:
while True:
msg = self.receive_msg(client, 5)
msglen, serialization = struct.unpack('>Lc', msg)
if self.verbose:
print "Msglen: " + str(msglen) + " is_serialized: " + str(serialization != 'n')
msg = self.receive_msg(client, msglen)
if serialization != 'n':
try:
if serialization == 'd': # dill serialization
msg = dill.loads(msg)
elif serialization == 'j': # json serialization
msg = json.loads(msg)
else:
print "Unrecognized serialization type: %s" % serialization
continue
except (UnicodeDecodeError, ValueError) as e:
print "Error de-serializing message: %s \n %s" % (
msg, traceback.format_exc(e))
continue
self.queue.put(msg)
except RuntimeError:
if self.verbose:
print "Client socket: " + str(address) + " closed"
def receive_msg(self, client, msglen):
msg = ''
while len(msg) < msglen:
chunk = client.recv(msglen-len(msg))
if chunk == '':
raise RuntimeError("socket connection broken")
msg = msg + chunk
return msg
def close(self):
self.sock.close()
class clientsocket:
"""A client socket for sending messages"""
def __init__(self, serialization='json', verbose=False):
""" `serialization` specifies the type of serialization to use for
non-str messages. Supported formats:
* 'json' uses the json module. Cross-language support. (default)
* 'dill' uses the dill pickle module. Python only.
"""
self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
if serialization != 'json' and serialization != 'dill':
raise ValueError("Unsupported serialization type: %s" % serialization)
self.serialization = serialization
self.verbose = verbose
def connect(self, host, port):
if self.verbose: print "Connecting to: %s:%i" % (host, port)
self.sock.connect((host, port))
def send(self, msg):
"""
Sends an arbitrary python object to the connected socket. Serializes
using dill if not str, and prepends msg len (4-bytes) and
serialization type (1-byte).
"""
#if input not string, serialize to string
if type(msg) is not str:
if self.serialization == 'dill':
msg = dill.dumps(msg)
serialization = 'd'
elif self.serialization == 'json':
msg = json.dumps(msg)
serialization = 'j'
else:
raise ValueError("Unsupported serialization type set: %s" % serialization)
else:
serialization = 'n'
if self.verbose: print "Sending message with serialization %s" % serialization
#prepend with message length
msg = struct.pack('>Lc', len(msg), serialization) + msg
totalsent = 0
while totalsent < len(msg):
sent = self.sock.send(msg[totalsent:])
if sent == 0:
raise RuntimeError("socket connection broken")
totalsent = totalsent + sent
def close(self):
self.sock.close()
if __name__ == '__main__':
import sys
#Just for testing
if sys.argv[1] == 's':
sock = serversocket(verbose=True)
sock.start_accepting()
raw_input("Press enter to exit...")
sock.close()
elif sys.argv[1] == 'c':
host = raw_input("Enter the host name:\n")
port = raw_input("Enter the port:\n")
serialization = raw_input("Enter the serialization type (default: 'json'):\n")
if serialization == '':
serialization = 'json'
sock = clientsocket(serialization=serialization)
sock.connect(host, int(port))
msg = None
# some predefined messages
tuple_msg = ('hello','world')
list_msg = ['hello','world']
dict_msg = {'hello':'world'}
def function_msg(x): return x
# read user input
while msg != "quit":
msg = raw_input("Enter a message to send:\n")
if msg == 'tuple':
sock.send(tuple_msg)
elif msg == 'list':
sock.send(list_msg)
elif msg == 'dict':
sock.send(dict_msg)
elif msg == 'function':
sock.send(function_msg)
else:
sock.send(msg)
sock.close()