1
|
1 |
|
import logging |
2
|
1 |
|
import os |
3
|
1 |
|
from concurrent.futures import Future |
4
|
1 |
|
import functools |
5
|
1 |
|
import threading |
6
|
|
|
|
7
|
1 |
|
try: |
8
|
|
|
# we prefer to use bundles asyncio version, otherwise fallback to trollius |
9
|
1 |
|
import asyncio |
10
|
|
|
except ImportError: |
11
|
|
|
import trollius as asyncio |
12
|
|
|
|
13
|
|
|
|
14
|
1 |
|
class UAError(RuntimeError): |
15
|
1 |
|
pass |
16
|
|
|
|
17
|
|
|
|
18
|
1 |
|
class ServiceError(UAError): |
19
|
1 |
|
def __init__(self, code): |
20
|
1 |
|
super(ServiceError, self).__init__('UA Service Error') |
21
|
1 |
|
self.code = code |
22
|
|
|
|
23
|
|
|
|
24
|
1 |
|
class NotEnoughData(UAError): |
25
|
1 |
|
pass |
26
|
|
|
|
27
|
|
|
|
28
|
1 |
|
class SocketClosedException(UAError): |
29
|
1 |
|
pass |
30
|
|
|
|
31
|
|
|
|
32
|
1 |
|
class Buffer(object): |
33
|
|
|
|
34
|
|
|
""" |
35
|
|
|
alternative to io.BytesIO making debug easier |
36
|
|
|
and added a few conveniance methods |
37
|
|
|
""" |
38
|
|
|
|
39
|
1 |
|
def __init__(self, data, start_pos=0, size=-1): |
40
|
|
|
# self.logger = logging.getLogger(__name__) |
41
|
1 |
|
self._data = data |
42
|
1 |
|
self._cur_pos = start_pos |
43
|
1 |
|
if size == -1: |
44
|
1 |
|
size = len(data) - start_pos |
45
|
1 |
|
self._size = size |
46
|
|
|
|
47
|
1 |
|
def __str__(self): |
48
|
|
|
return "Buffer(size:{}, data:{})".format( |
49
|
|
|
self._size, |
50
|
|
|
self._data[self._cur_pos:self._cur_pos + self._size]) |
51
|
1 |
|
__repr__ = __str__ |
52
|
|
|
|
53
|
1 |
|
def __len__(self): |
54
|
1 |
|
return self._size |
55
|
|
|
|
56
|
1 |
|
def read(self, size): |
57
|
|
|
""" |
58
|
|
|
read and pop number of bytes for buffer |
59
|
|
|
""" |
60
|
1 |
|
if size > self._size: |
61
|
|
|
raise NotEnoughData("Not enough data left in buffer, request for {}, we have {}".format(size, self)) |
62
|
|
|
# self.logger.debug("Request for %s bytes, from %s", size, self) |
63
|
1 |
|
self._size -= size |
64
|
1 |
|
pos = self._cur_pos |
65
|
1 |
|
self._cur_pos += size |
66
|
1 |
|
data = self._data[pos:self._cur_pos] |
67
|
|
|
# self.logger.debug("Returning: %s ", data) |
68
|
1 |
|
return data |
69
|
|
|
|
70
|
1 |
|
def copy(self, size=-1): |
71
|
|
|
""" |
72
|
|
|
return a shadow copy, optionnaly only copy 'size' bytes |
73
|
|
|
""" |
74
|
1 |
|
if size == -1 or size > self._size: |
75
|
1 |
|
size = self._size |
76
|
1 |
|
return Buffer(self._data, self._cur_pos, size) |
77
|
|
|
|
78
|
1 |
|
def skip(self, size): |
79
|
|
|
""" |
80
|
|
|
skip size bytes in buffer |
81
|
|
|
""" |
82
|
1 |
|
if size > self._size: |
83
|
|
|
raise NotEnoughData("Not enough data left in buffer, request for {}, we have {}".format(size, self)) |
84
|
1 |
|
self._size -= size |
85
|
1 |
|
self._cur_pos += size |
86
|
|
|
|
87
|
|
|
|
88
|
1 |
|
class SocketWrapper(object): |
89
|
|
|
""" |
90
|
|
|
wrapper to make it possible to have same api for |
91
|
|
|
normal sockets, socket from asyncio, StringIO, etc.... |
92
|
|
|
""" |
93
|
1 |
|
def __init__(self, sock): |
94
|
1 |
|
self.socket = sock |
95
|
|
|
|
96
|
1 |
|
def read(self, size): |
97
|
|
|
""" |
98
|
|
|
Receive up to size bytes from socket |
99
|
|
|
""" |
100
|
1 |
|
data = b'' |
101
|
1 |
|
while size > 0: |
102
|
1 |
|
chunk = self.socket.recv(size) |
103
|
1 |
|
if not chunk: |
104
|
1 |
|
raise SocketClosedException("Server socket has closed") |
105
|
1 |
|
data += chunk |
106
|
1 |
|
size -= len(chunk) |
107
|
1 |
|
return data |
108
|
|
|
|
109
|
1 |
|
def write(self, data): |
110
|
1 |
|
self.socket.sendall(data) |
111
|
|
|
|
112
|
|
|
|
113
|
1 |
|
def create_nonce(size=32): |
114
|
1 |
|
return os.urandom(size) |
115
|
|
|
|
116
|
|
|
|
117
|
1 |
|
class ThreadLoop(threading.Thread): |
118
|
|
|
""" |
119
|
|
|
run an asyncio loop in a thread |
120
|
|
|
""" |
121
|
|
|
|
122
|
1 |
|
def __init__(self): |
123
|
1 |
|
threading.Thread.__init__(self) |
124
|
1 |
|
self.logger = logging.getLogger(__name__) |
125
|
1 |
|
self.loop = None |
126
|
1 |
|
self._cond = threading.Condition() |
127
|
|
|
|
128
|
1 |
|
def start(self): |
129
|
1 |
|
with self._cond: |
130
|
1 |
|
threading.Thread.start(self) |
131
|
1 |
|
self._cond.wait() |
132
|
|
|
|
133
|
1 |
|
def run(self): |
134
|
1 |
|
self.logger.debug("Starting subscription thread") |
135
|
1 |
|
self.loop = asyncio.new_event_loop() |
136
|
1 |
|
asyncio.set_event_loop(self.loop) |
137
|
1 |
|
with self._cond: |
138
|
1 |
|
self._cond.notify_all() |
139
|
1 |
|
self.loop.run_forever() |
140
|
1 |
|
self.logger.debug("subscription thread ended") |
141
|
|
|
|
142
|
1 |
|
def create_server(self, proto, hostname, port): |
143
|
1 |
|
return self.loop.create_server(proto, hostname, port) |
144
|
|
|
|
145
|
1 |
|
def stop(self): |
146
|
|
|
""" |
147
|
|
|
stop subscription loop, thus the subscription thread |
148
|
|
|
""" |
149
|
1 |
|
self.loop.call_soon_threadsafe(self.loop.stop) |
150
|
|
|
|
151
|
1 |
|
def call_soon(self, callback): |
152
|
1 |
|
self.loop.call_soon_threadsafe(callback) |
153
|
|
|
|
154
|
1 |
|
def call_later(self, delay, callback): |
155
|
|
|
""" |
156
|
|
|
threadsafe call_later from asyncio |
157
|
|
|
""" |
158
|
1 |
|
p = functools.partial(self.loop.call_later, delay, callback) |
159
|
1 |
|
self.loop.call_soon_threadsafe(p) |
160
|
|
|
|
161
|
1 |
|
def _create_task(self, future, coro, cb=None): |
162
|
|
|
#task = self.loop.create_task(coro) |
163
|
1 |
|
task = asyncio.async(coro, loop=self.loop) |
164
|
1 |
|
if cb: |
165
|
1 |
|
task.add_done_callback(cb) |
166
|
1 |
|
future.set_result(task) |
167
|
|
|
|
168
|
1 |
|
def create_task(self, coro, cb=None): |
169
|
|
|
""" |
170
|
|
|
threadsafe create_task from asyncio |
171
|
|
|
""" |
172
|
1 |
|
future = Future() |
173
|
1 |
|
p = functools.partial(self._create_task, future, coro, cb) |
174
|
1 |
|
self.loop.call_soon_threadsafe(p) |
175
|
1 |
|
return future.result() |
176
|
|
|
|
177
|
1 |
|
def run_coro_and_wait(self, coro): |
178
|
1 |
|
cond = threading.Condition() |
179
|
1 |
|
def cb(_): |
180
|
1 |
|
with cond: |
181
|
1 |
|
cond.notify_all() |
182
|
1 |
|
with cond: |
183
|
1 |
|
task = self.create_task(coro, cb) |
184
|
1 |
|
cond.wait() |
185
|
1 |
|
return task.result() |
186
|
|
|
|
187
|
1 |
|
def _run_until_complete(self, future, coro): |
188
|
|
|
task = self.loop.run_until_complete(coro) |
189
|
|
|
future.set_result(task) |
190
|
|
|
|
191
|
1 |
|
def run_until_complete(self, coro): |
192
|
|
|
""" |
193
|
|
|
threadsafe run_until_completed from asyncio |
194
|
|
|
""" |
195
|
|
|
future = Future() |
196
|
|
|
p = functools.partial(self._run_until_complete, future, coro) |
197
|
|
|
self.loop.call_soon_threadsafe(p) |
198
|
|
|
return future.result() |
199
|
|
|
|
200
|
|
|
|
201
|
|
|
|
202
|
|
|
|