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
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
|
import aiohttp
import aiohttp.web
import asyncio
import collections
import concurrent.futures
import json
import logging
import signal
logging.basicConfig(level = logging.DEBUG, format = '{asctime} {levelname} {message}', style = '{')
class MessageQueue:
# An object holding onto the messages received from nodeping
# This is effectively a reimplementation of parts of asyncio.Queue with some specific additional code.
# Unfortunately, asyncio.Queue's extensibility (_init, _put, and _get methods) is undocumented, so I don't want to rely on that.
# Differences to asyncio.Queue include:
# - No maxsize
# - No put coroutine (not necessary since the queue can never be full)
# - Only one concurrent getter
# - putleft_nowait to put to the front of the queue (so that the IRC client can put a message back when delivery fails)
def __init__(self):
self._getter = None # None | asyncio.Future
self._queue = collections.deque()
async def get(self):
if self._getter is not None:
raise RuntimeError('Cannot get concurrently')
if len(self._queue) == 0:
self._getter = asyncio.get_running_loop().create_future()
logging.debug('Awaiting getter')
try:
await self._getter
except asyncio.CancelledError:
logging.debug('Cancelled getter')
self._getter = None
raise
logging.debug('Awaited getter')
self._getter = None
# For testing the cancellation/putting back onto the queue
#logging.debug('Delaying message queue get')
#await asyncio.sleep(3)
#logging.debug('Done delaying')
return self.get_nowait()
def get_nowait(self):
if len(self._queue) == 0:
raise asyncio.QueueEmpty
return self._queue.popleft()
def put_nowait(self, item):
self._queue.append(item)
if self._getter is not None:
self._getter.set_result(None)
def putleft_nowait(self, item):
self._queue.appendleft(item)
if self._getter is not None:
self._getter.set_result(None)
def qsize(self):
return len(self._queue)
class IRCClientProtocol(asyncio.Protocol):
def __init__(self, messageQueue, stopEvent, loop):
logging.debug(f'Protocol init {id(self)}: {messageQueue} {id(messageQueue)}, {stopEvent}, {loop}')
self.messageQueue = messageQueue
self.stopEvent = stopEvent
self.loop = loop
self.buffer = b''
self.connected = False
def send(self, data):
logging.info(f'Send: {data!r}')
self.transport.write(data + b'\r\n')
def connection_made(self, transport):
logging.info('Connected')
self.transport = transport
self.connected = True
self.send(b'NICK npbot')
self.send(b'USER npbot npbot npbot :I am a bot.')
self.send(b'JOIN #nodeping')
asyncio.create_task(self.send_messages())
async def _get_message(self):
logging.debug(f'Message queue {id(self.messageQueue)} length: {self.messageQueue.qsize()}')
messageFuture = asyncio.create_task(self.messageQueue.get())
done, pending = await asyncio.wait((messageFuture, self.stopEvent.wait()), return_when = concurrent.futures.FIRST_COMPLETED)
if self.stopEvent.is_set():
if messageFuture in pending:
logging.debug('Cancelling messageFuture')
messageFuture.cancel()
try:
await messageFuture
except asyncio.CancelledError:
logging.debug('Cancelled messageFuture')
pass
else:
# messageFuture is already done but we're stopping, so put the result back onto the queue
self.messageQueue.putleft_nowait(messageFuture.result())
return None
assert messageFuture in done, 'Invalid state: messageFuture not in done futures'
return messageFuture.result()
async def send_messages(self):
while self.connected:
logging.debug(f'{id(self)}: trying to get a message')
message = await self._get_message()
logging.debug(f'{id(self)}: got message: {message!r}')
if message is None:
break
self.send(b'PRIVMSG #nodeping :' + message.encode('utf-8'))
#TODO self.messageQueue.putleft_nowait if delivery fails
await asyncio.sleep(1) # Rate limit
def data_received(self, data):
logging.debug(f'Data received: {data!r}')
# Split received data on CRLF. If there's any data left in the buffer, prepend it to the first message and process that.
# Then, process all messages except the last one (since data might not end on a CRLF) and keep the remainder in the buffer.
# If data does end with CRLF, all messages will have been processed and the buffer will be empty again.
messages = data.split(b'\r\n')
if self.buffer:
self.message_received(self.buffer + messages[0])
messages = messages[1:]
for message in messages[:-1]:
self.message_received(message)
self.buffer = messages[-1]
def message_received(self, message):
logging.info(f'Message received: {message!r}')
if message.startswith(b'PING '):
self.send(b'PONG ' + message[5:])
def connection_lost(self, exc):
logging.info('The server closed the connection')
self.connected = False
self.stopEvent.set()
class WebServer:
def __init__(self, messageQueue):
self.messageQueue = messageQueue
self._app = aiohttp.web.Application()
self._app.add_routes([aiohttp.web.post('/nodeping', self.nodeping_post)])
async def run(self, stopEvent):
runner = aiohttp.web.AppRunner(self._app)
await runner.setup()
site = aiohttp.web.TCPSite(runner, '127.0.0.1', 8080)
await site.start()
await stopEvent.wait()
await runner.cleanup()
async def nodeping_post(self, request):
logging.info(f'Received request with data: {await request.read()!r}')
authHeader = request.headers.get('Authorization')
if not authHeader or authHeader != 'Basic YXJjaGl2ZXRlYW06aXNvbmZpcmU=': #TODO move out of source code
return aiohttp.web.HTTPForbidden()
try:
data = await request.json()
except (aiohttp.ContentTypeError, json.JSONDecodeError) as e:
logging.error(f'Received invalid data: {await request.read()!r}')
return aiohttp.web.HTTPBadRequest()
if 'message' not in data:
logging.error(f'Received invalid data: {await request.read()!r}')
return aiohttp.web.HTTPBadRequest()
if '\r' in data['message'] or '\n' in data['message']:
logging.error(f'Received invalid data: {await request.read()!r}')
return aiohttp.web.HTTPBadRequest()
logging.debug(f'Putting to message queue {id(self.messageQueue)}')
self.messageQueue.put_nowait(data['message'])
return aiohttp.web.HTTPOk()
async def run_irc(loop, messageQueue, sigintEvent):
stopEvent = asyncio.Event()
while True:
stopEvent.clear()
try:
# transport, protocol = await loop.create_connection(lambda: IRCClientProtocol(messageQueue, stopEvent, loop), 'irc.hackint.org', 6697, ssl = True)
transport, protocol = await loop.create_connection(lambda: IRCClientProtocol(messageQueue, stopEvent, loop), '127.0.0.1', 8888)
try:
await asyncio.wait((stopEvent.wait(), sigintEvent.wait()), return_when = concurrent.futures.FIRST_COMPLETED)
finally:
transport.close()
except (ConnectionRefusedError, asyncio.TimeoutError) as e:
logging.error(str(e))
await asyncio.wait((asyncio.sleep(5), sigintEvent.wait()), return_when = concurrent.futures.FIRST_COMPLETED)
if sigintEvent.is_set():
break
async def run_webserver(loop, messageQueue, sigintEvent):
server = WebServer(messageQueue)
await server.run(sigintEvent)
async def main():
loop = asyncio.get_running_loop()
messageQueue = MessageQueue()
sigintEvent = asyncio.Event()
def sigint_callback():
logging.info('Got SIGINT')
nonlocal sigintEvent
sigintEvent.set()
loop.add_signal_handler(signal.SIGINT, sigint_callback)
await asyncio.gather(run_irc(loop, messageQueue, sigintEvent), run_webserver(loop, messageQueue, sigintEvent))
asyncio.run(main())
|
