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
|
from ws4py import configure_logger
from ws4py.websocket import WebSocket
from ws4py.messaging import Message, TextMessage
logger = configure_logger(stdout=True)
class DataSource(object):
def __init__(self):
self.frames = set()
self.frame = None
self.remaining_bytes = None
def setblocking(self, flag):
pass
def feed(self, message):
if isinstance(message, Message):
message = message.single(mask=True)
else:
message = TextMessage(message).single(mask=True)
self.frames.add(message)
def recv(self, size):
if not self.frame:
if not self.frames:
return b''
self.frame = self.frames.pop()
self.remaining_bytes = self.frame
current_bytes = self.remaining_bytes[:size]
self.remaining_bytes = self.remaining_bytes[size:]
if self.remaining_bytes is b'':
self.frame = None
self.remaining_bytes = None
return current_bytes
class LogWebSocket(WebSocket):
def opened(self):
logger.info("WebSocket now ready")
def closed(self, code=1000, reason="Burp, done!"):
logger.info("Terminated with reason '%s'" % reason)
def received_message(self, m):
logger.info("Received message: %s" % m)
if __name__ == '__main__':
source = DataSource()
ws = LogWebSocket(sock=source)
source.feed(u'hello there')
source.feed(u'a bit more')
ws.run()
|
