async def main():
pictures_queue = Queue()
workers_count = 300
connection = {
'user': '******', # input your postgres username
'database': 'your database name', # input your database name here
'host': '127.0.0.1', # change your host if it's not local
'password': '******' # input your password for this database
}
dsn = 'postgresql://{user}:{password}@{host}/{database}'.format(**connection)
engine = create_engine(dsn)
result = engine.execute('''select picture from "your_table_name"''')
res_list = []
for row in result:
clean_jpg = row['picture'].split("\n")
for i in clean_jpg:
res_list.append(i)
print(len(res_list))
for pic in res_list:
pictures_queue.put_nowait(pic)
session = AsyncHTMLSession()
tasks = []
for num in range(workers_count):
task = worker(pictures_queue, num, session)
tasks.append(task)
await asyncio.gather(*tasks)
class Database:
def __init__(self, database, loop=None):
self.database = database
self.loop = loop if loop else asyncio.get_event_loop()
self.connection_semaphore = Queue(maxsize=1)
self.connection_semaphore.put_nowait({})
self.__connect_ref = None
async def __acquire_connection(self):
await self.connection_semaphore.get()
return sqlite3.connect(self.database, check_same_thread=False)
async def __call__(self, *args, **kwargs):
db = await self.__acquire_connection()
return Connection(db, self)
async def __aenter__(self):
db = await self.__acquire_connection()
conn = Connection(db, self)
self.__connect_ref = conn
return conn
async def __aexit__(self, exc_type, exc_val, exc_tb):
if self.__connect_ref:
await self.__connect_ref.close()
class TaskPool(object):
def __init__(self, num_workers):
self.loop = asyncio.get_event_loop()
self.tasks = Queue(loop=self.loop)
self.workers = []
for _ in range(num_workers):
worker = asyncio.create_task(self.worker())
self.workers.append(worker)
async def worker(self):
while True:
future, task = await self.tasks.get()
if task == "TERMINATOR":
break
result = await asyncio.wait_for(task, None, loop=self.loop)
future.set_result(result)
def submit(self, task):
future = asyncio.Future(loop=self.loop)
self.tasks.put_nowait((future, task))
return future
async def close(self):
for _ in self.workers:
self.tasks.put_nowait((None, "TERMINATOR"))
await asyncio.gather(*self.workers, loop=self.loop)
class QueueMailBox(MailBox):
def __init__(self, name, *args, **kwargs):
super(QueueMailBox, self).__init__(*args, **kwargs)
self._name = name
self._queue = Queue()
self._ready = True
async def prepare(self):
self._ready = True
async def put(self, msg=None):
if await self.policy():
self._queue.put_nowait(msg)
async def size(self):
return self._queue.qsize()
async def empty(self):
return self._queue.empty()
async def get(self):
result = None
result = await self._queue.get()
return result
async def policy(self):
mem_percent = psutil.virtual_memory().percent
if mem_percent > 80:
logger.warning("memory usage is gt than 80")
return True
class TaskPool(object):
def __init__(self, loop, num_workers):
self.loop = loop
self.tasks = Queue(loop=self.loop)
self.workers = []
for _ in range(num_workers):
worker = asyncio.ensure_future(self.worker(), loop=self.loop)
self.workers.append(worker)
async def worker(self):
while True:
future, task = await self.tasks.get()
if task is TERMINATOR:
break
result = await asyncio.wait_for(task, None, loop=self.loop)
future.set_result(result)
def submit(self, task):
future = asyncio.Future(loop=self.loop)
self.tasks.put_nowait((future, task))
return future
async def join(self):
for _ in self.workers:
self.tasks.put_nowait((None, TERMINATOR))
await asyncio.gather(*self.workers, loop=self.loop)
class ChannelHandler(logging.Handler):
def __init__(self, bot: discord.Client, cog, cog_name: str,
channel: discord.Channel, *args, **kwargs):
self.cog = cog
self.cog_name = cog_name
self.queue = Queue()
self.channel = channel
self.bot = bot
super(ChannelHandler, self).__init__(*args, **kwargs)
def emit(self, record):
if self.channel:
self.queue.put_nowait(record)
async def update_task(self):
while self.cog == self.bot.get_cog(self.cog_name):
if not self.channel:
await asyncio.sleep(1)
else:
record = await self.queue.get()
if self.cog != self.bot.get_cog(self.cog_name):
return
msg = self.format(record)
await self.bot.send_message(content=msg,
destination=self.channel)
class IRCClientProtocol(asyncio.Protocol):
"""Low-level protocol that speaks the client end of IRC.
This isn't responsible for very much besides the barest minimum definition
of an IRC client: connecting and responding to PING.
You probably want `read_message`, or the higher-level client class.
"""
def __init__(self, loop, nick, password, charset="utf8"):
self.nick = nick
self.password = password
self.charset = charset
self.buf = b""
self.message_queue = Queue(loop=loop)
self.registered = False
def connection_made(self, transport):
self.transport = transport
if self.password:
self.send_message("PASS", self.password)
self.send_message("NICK", self.nick)
self.send_message("USER", "dywypi", "-", "-", "dywypi Python IRC bot")
def data_received(self, data):
data = self.buf + data
while True:
raw_message, delim, data = data.partition(b"\r\n")
if not delim:
# Incomplete message; stop here and wait for more
self.buf = raw_message
return
# TODO valerr
message = IRCMessage.parse(raw_message.decode(self.charset))
logger.debug("recv: %r", message)
self.handle_message(message)
def handle_message(self, message):
if message.command == "PING":
self.send_message("PONG", message.args[-1])
elif message.command == "RPL_WELCOME":
# 001, first thing sent after registration
if not self.registered:
self.registered = True
self.message_queue.put_nowait(message)
def send_message(self, command, *args):
message = IRCMessage(command, *args)
logger.debug("sent: %r", message)
self.transport.write(message.render().encode(self.charset) + b"\r\n")
@asyncio.coroutine
def read_message(self):
return (yield from self.message_queue.get())
async def _serve(self, websocket, path):
import websockets
connection_id = await websocket.recv()
with self.queue_lock:
#if connection_id in self._preclosed_connections:
# self._preclosed_connections[connection_id].set()
# return
pmqueue = self._pending_message_queues.get(connection_id, None)
myqueue = Queue()
if pmqueue is not None:
events = []
while 1:
try:
e = pmqueue.get_nowait()
except QueueEmpty:
break
events.append(e)
if len(events) > self.CACHE_EVENTS_FIRST + self.CACHE_EVENTS_LAST:
events = events[:self.CACHE_EVENTS_FIRST] + \
events[-self.CACHE_EVENTS_LAST:]
for enr, e in enumerate(events):
swallow = False
if e.get("type", None) == "var":
varname = e.get("var", None)
if varname is not None:
for e2 in events[enr+1:]:
if e2.get("type", None) != "var":
continue
if e2.get("var", None) != varname:
continue
swallow = True
break
if swallow:
continue
myqueue.put_nowait(e)
pmqueue.put_nowait(e) #put the events back
if connection_id not in self._message_queue_lists:
self._message_queue_lists[connection_id] = []
self._message_queue_lists[connection_id].append(myqueue)
while True:
#print("WAIT")
try:
message = await myqueue.get()
except Exception:
break
message = json.dumps(message)
#print("SEND?", message)
if message is None: #terminating message
break
try:
#print("SEND", message)
await websocket.send(message)
except websockets.exceptions.ConnectionClosed:
break
with self.queue_lock:
self._message_queue_lists[connection_id].remove(myqueue)
class Core(object):
def __init__(self, bot):
self.bot = bot
self.timeout = int(self.bot.config.get('timeout'))
self.ping_queue = Queue(loop=bot.loop)
def connection_made(self):
self.bot.loop.call_later(self.timeout, self.check_ping)
self.ping_queue.put_nowait(self.bot.loop.time())
def check_ping(self): # pragma: no cover
# check if we received a ping
# reconnect if queue is empty
self.bot.log.debug(
'Ping queue size: {}'.format(self.ping_queue.qsize()))
if self.ping_queue.empty():
self.bot.loop.call_soon(self.bot.protocol.transport.close)
else:
self.bot.loop.call_later(self.timeout, self.check_ping)
while not self.ping_queue.empty():
self.ping_queue.get_nowait()
@event(rfc.PING)
def pong(self, data):
"""PING reply"""
self.ping_queue.put_nowait(self.bot.loop.time())
self.bot.send('PONG ' + data)
@event(rfc.NEW_NICK)
def recompile(self, nick=None, new_nick=None, **kw):
"""recompile regexp on new nick"""
if self.bot.nick == nick.nick:
self.bot.config['nick'] = new_nick
self.bot.recompile()
@event(rfc.ERR_NICK)
def badnick(self, me=None, nick=None, **kw):
"""Use alt nick on nick error"""
if me == '*':
self.bot.set_nick(self.bot.nick + '_')
self.bot.log.debug('Trying to regain nickname in 30s...')
self.bot.loop.call_later(30, self.bot.set_nick, self.bot.original_nick)
@event(rfc.RPL_ENDOFMOTD)
def autojoin(self, **kw):
"""autojoin at the end of MOTD"""
self.bot.config['nick'] = kw['me']
self.bot.recompile()
channels = utils.as_list(self.bot.config.get('autojoins', []))
for channel in channels:
channel = utils.as_channel(channel)
self.bot.log.info('Trying to join %s', channel)
self.bot.join(channel)
async def searchWorker(name: str, searchQueue: Queue, queueProgress: tqdm):
while True:
# Get a "work item" out of the queue.
data: QueueObj = await searchQueue.get()
query = data['query']
url = queryToUrl(query)
print(Fore.MAGENTA + f'{name} got url: {url} to work on' + Fore.RESET)
start = time.time()
custom_timeout = aiohttp.ClientTimeout(total=60) # type: ignore
# custom_timeout.total = 2*60
async with aiohttp.ClientSession(timeout=custom_timeout) as session:
# nonlocal results
results = []
try:
results = await getSearchQueryResults(session, url)
except asyncio.TimeoutError as e:
logging.exception(
Fore.RED +
f'Exception raised by worker = {name}; error = {e}' +
Fore.RESET)
except Exception as e: # pylint: disable=broad-except
logging.exception(
Fore.RED +
f'Exception raised by worker = {name}; error = {e}' +
Fore.RESET)
print(
Fore.YELLOW +
f'{name}\'s work for url {url} done; time taken {time.time() - start} seconds'
+ Fore.RESET)
if (len(results) > 0 and len(results[0]['results']) > 0
and len(results[0]['results'][query]) > 0):
print(f'query: \'{query}\' fetched with data'.encode(
encoding='utf-8'))
allResults = fromSearchData(results, query)
searchQueue.task_done()
onlyTitles = [item['title'] for item in allResults]
relevantIndices = getRelevantTitleIndices(onlyTitles)
relevantSearchData = [
allResults[index] for index in relevantIndices
]
bookRow = booksDf[booksDf['id'] == data['id']].copy()
bookRow['searchLinks'] = [relevantSearchData]
bookRow['query'] = query
saveLinks(bookRow)
queueProgress.update(1)
# searchResults.append({
# query: results
# })
else:
print(f'query: \'{query}\' unsuccessful, adding to queue again'.
encode(encoding='utf-8'))
searchQueue.task_done()
searchQueue.put_nowait(data)
class Port(object):
def __init__(self, tag="data", maxsize=1, name=None, loop=None):
loop = loop if loop is not None else asyncio.get_event_loop()
self.loop = loop
self.name = name if name is not None else str(uuid1())
self._queue = Queue(maxsize, loop=self.loop)
self.default_value = None
self.default_value_set = False
self.connected = False
self.belong_to_block = None
self.data_tag = tag
def set_default_value(self, value):
if not isinstance(value, Payload):
raise Exception("value should be Payload type")
self.default_value = value
self.default_value_set = True
async def get(self):
if self.default_value_set:
if self._queue.empty():
return self.default_value, self.default_value[self.data_tag]
payload = await self._queue.get()
return payload, payload[self.data_tag]
def get_nowait(self):
if self.default_value_set:
if self._queue.empty():
return self.default_value, self.default_value[self.data_tag]
payload = self._queue.get_nowait()
return payload, payload[self.data_tag]
async def put(self, payload, item):
if self.connected:
payload[self.data_tag] = item
await self._queue.put(payload)
def put_nowait(self, payload, item):
if self.connected:
payload[self.data_tag] = item
self._queue.put_nowait(payload)
def empty(self):
return self._queue.empty()
def full(self):
return self._queue.full()
def set_buffer_size(self, maxsize):
self._queue = Queue(maxsize, loop=self.loop)
async def test_commit_concurrency(aconn):
# Check the condition reported in psycopg2#103
# Because of bad status check, we commit even when a commit is already on
# its way. We can detect this condition by the warnings.
notices = Queue()
aconn.add_notice_handler(
lambda diag: notices.put_nowait(diag.message_primary))
stop = False
async def committer():
nonlocal stop
while not stop:
await aconn.commit()
await asyncio.sleep(0) # Allow the other worker to work
async def runner():
nonlocal stop
cur = aconn.cursor()
for i in range(1000):
await cur.execute("select %s;", (i, ))
await aconn.commit()
# Stop the committer thread
stop = True
await asyncio.gather(committer(), runner())
assert notices.empty(), "%d notices raised" % notices.qsize()
class FoundComponentIterator:
def __init__(self, loop, components_file_request_list):
self.loop = loop
self.pending_tasks = components_file_request_list
self._to_remove = None
self.done = Queue(loop=self.loop)
self.add_done_callback()
def add_done_callback(self):
for task in self.pending_tasks:
task.add_done_callback(self.on_completion)
def on_completion(self, task):
self.pending_tasks.remove(task)
self.done.put_nowait(task)
async def cancel_pending_tasks(self):
for task in self.pending_tasks:
task.cancel()
if len(self.pending_tasks):
await asyncio.wait(self.pending_tasks)
while not self.done.empty():
task = self.done.get_nowait()
try:
task.result()
except:
pass
def __aiter__(self):
return self
async def __anext__(self):
try:
while len(self.pending_tasks) > 0 or not self.done.empty():
try:
future = await self.done.get()
component_key, fetched_files = await future
self._to_remove = future
if len(fetched_files) > 0:
return {'key': component_key, 'files': fetched_files}
except (RejectRequest, StopRequest):
# Not fatal at all, just one of many
pass
except OfflineHostException:
await self.cancel_pending_tasks()
raise
raise StopAsyncIteration
async def main():
category_queue = Queue()
workers_count = 100
category_base_url = 'https://rozetka.com.ua/ua/notebooks/c80004/page={}'
for i in range(1, 100):
cat_url = category_base_url.format(i)
category_queue.put_nowait(cat_url)
session = AsyncHTMLSession()
tasks = []
for num in range(workers_count):
task = worker(category_queue, num, session)
tasks.append(task)
await asyncio.gather(*tasks)
async def main():
category_queue = Queue()
workers_count = 100
with open('notebook_links.txt', 'r', encoding='utf-8') as f:
links = f.read().split('\n')
for link in links:
category_queue.put_nowait(link)
session = AsyncHTMLSession()
engine1 = await create_engine(**connection)
tasks = []
for num in range(workers_count):
task = worker(category_queue, num, session, engine1)
tasks.append(task)
await asyncio.gather(*tasks)
class IrcConnection(asyncio.Protocol):
"""asyncio protocol to handle an irc connection"""
def connection_made(self, transport):
self.transport = transport
self.closed = False
self.queue = Queue()
def data_received(self, data):
encoding = getattr(self, 'encoding', 'ascii')
data = data.decode(encoding, 'ignore')
if not self.queue.empty():
data = self.queue.get_nowait() + data
lines = data.split('\r\n')
self.queue.put_nowait(lines.pop(-1))
for line in lines:
self.factory.dispatch(line)
def write(self, data):
if data is not None:
if isinstance(data, text_type):
data = data.encode(self.encoding)
if not data.endswith(b'\r\n'):
data = data + b'\r\n'
self.transport.write(data)
def connection_lost(self, exc): # pragma: no cover
self.factory.log.critical('connection lost (%s): %r',
id(self.transport),
exc)
self.factory.notify('connection_lost')
if not self.closed:
self.close()
# wait a few before reconnect
self.factory.loop.call_later(
2, self.factory.create_connection, self.__class__)
def close(self): # pragma: no cover
if not self.closed:
self.factory.log.critical('closing old transport (%r)',
id(self.transport))
try:
self.transport.close()
finally:
self.closed = True
async def main():
urls_queue = Queue()
workers_count = 100
pages = 97
category_base_url = 'https://www.searchengines.ru/category/seo/page/{}'
for i in range(1, pages):
cat_url = category_base_url.format(i)
urls_queue.put_nowait(cat_url)
session = AsyncHTMLSession()
engine = await create_engine(**connection)
tasks = []
for num in range(workers_count):
task = worker(urls_queue, num, session, engine)
tasks.append(task)
await asyncio.gather(*tasks)
class IrcConnection(asyncio.Protocol):
"""asyncio protocol to handle an irc connection"""
def connection_made(self, transport):
self.transport = transport
self.closed = False
self.queue = Queue()
def data_received(self, data):
encoding = getattr(self, 'encoding', 'ascii')
data = data.decode(encoding, 'ignore')
if not self.queue.empty():
data = self.queue.get_nowait() + data
lines = data.split('\r\n')
self.queue.put_nowait(lines.pop(-1))
for line in lines:
self.factory.dispatch(line)
def write(self, data):
if data is not None:
if isinstance(data, text_type):
data = data.encode(self.encoding)
if not data.endswith(b'\r\n'):
data = data + b'\r\n'
self.transport.write(data)
def connection_lost(self, exc): # pragma: no cover
self.factory.log.critical('connection lost (%s): %r',
id(self.transport), exc)
self.factory.notify('connection_lost')
if not self.closed:
self.close()
# wait a few before reconnect
self.factory.loop.call_later(2, self.factory.create_connection,
self.__class__)
def close(self): # pragma: no cover
if not self.closed:
self.factory.log.critical('closing old transport (%r)',
id(self.transport))
try:
self.transport.close()
finally:
self.closed = True
class LocalChat(BaseChat):
def __init__(self, bridge):
super().__init__(bridge)
self.queue = Queue()
def _create_channel(self, channel):
return channel
async def _send_msg(self, msg):
print(msg.to_dict())
print(msg.message)
async def _listen_to(self, channel):
loop = get_running_loop()
loop.add_reader(sys.stdin,
lambda: self.queue.put_nowait(sys.stdin.readline()))
while True:
msg = await self.queue.get()
self.handle_msg(
IncomingMessage(message=msg,
channel=channel,
date=datetime.now(tz=timezone.utc),
**_BASE_MSG))
async def dump1090Receiver(incomingQ: Queue,
db_worker_Q: Queue = None,
radar_Q=None,
websocket_Q=None):
"""
Receive a raw mesage from a Queue parse it sent it to the database_worker (if queue is not None) process the message otherwise)
:param incomingQ: incoming messaged from 1090
:param db_worker_Q: queue to send processed messages to db layer -- if not None
:param radar_Q: queue to send processed messages to radar -- if not None
:return:
"""
LOG.info("Starting")
try:
while True:
raw_sbs1_msg = await incomingQ.get()
#print(f"{raw_item}")
sbs1_msg = SBS1Message(raw_sbs1_msg)
if sbs1_msg.isValid:
#process(sbs1_msg)
result = process2(sbs1_msg)
if result is None:
continue
# higher level
new_aircraft = result.new_aircraft
position = any([
result.callsign, result.squawk, result.altitude,
result.ground_speed, result.track, result.lat, result.lon,
result.vertical_rate, result.on_ground
])
emergency = any([result.alert, result.emergency])
if any([new_aircraft, position, emergency]):
#print(F"[{sbs1_msg.messageType}{sbs1_msg.transmissionType}] New Aircraft: {new_aircraft} Position: {position} Emergency: {emergency}")
#dump_row(db2, sbs1_msg.icao24)
actionTypesList = []
if new_aircraft:
actionTypesList.append("new_aircraft")
if position:
actionTypesList.append("position")
if emergency:
actionTypesList.append("emergency")
incoming_action = INCOMING_ACTION_TYPE(
actionTypes=actionTypesList,
actionMsg=result.icao_record,
aircraftId=sbs1_msg.icao24)
STATS["incoming"] += 1
if STATS["incoming"] % 100 == 0:
LOG.info(F"Incoming {STATS['incoming']}")
if db_worker_Q:
db_worker_Q.put_nowait(incoming_action)
if radar_Q:
radar_Q.put_nowait(incoming_action)
if websocket_Q:
websocket_Q.put_nowait(incoming_action)
flush_cache(300)
else:
STATS["parsing_error"] += 1
LOG.warning(
F"Invalid sbs1_msg. Failed processing count: {STATS['parsing_error']}"
)
except CancelledError:
LOG.info("Cancelling")
except Exception:
LOG.exception("Any exception")
finally:
pass
LOG.info("Exiting")
class DywypiShell(UrwidProtocol):
"""Creates a Twisted-friendly urwid app that allows interacting with dywypi
via a shell.
"""
# TODO i don't think client.nick should really be part of the exposed
# interface; should be a .me returning a peer probably
# TODO for some reason this is the client even though the thing below is
# actually called a Client so we should figure this the f**k out
nick = 'dywypi'
def __init__(self, loop, network, *args, **kwargs):
super().__init__(loop, **kwargs)
self.event_queue = Queue(loop=self.loop)
self.network = network
self.me = Peer('dywypi', 'dywypi', 'localhost')
self.you = Peer('user', 'user', 'localhost')
def build_toplevel_widget(self):
self.pane = UnselectableListBox(urwid.SimpleListWalker([]))
prompt = FancyEdit('>>> ')
urwid.connect_signal(prompt, 'line_submitted', self.handle_line)
return urwid.Pile(
[
self.pane,
('flow', prompt),
],
focus_item=prompt,
)
def build_palette(self):
return [
('default', 'default', 'default'),
('logging-debug', 'dark gray', 'default'),
('logging-info', 'light gray', 'default'),
('logging-warning', 'light red', 'default'),
('logging-error', 'dark red', 'default'),
('logging-critical', 'light magenta', 'default'),
('shell-input', 'light gray', 'default'),
('bot-output', 'default', 'default'),
('bot-output-label', 'dark cyan', 'default'),
]
def unhandled_input(self, key):
# Try passing the key along to the listbox, so pgup/pgdn still work.
# Note that this is a Pile method specifically, and requires an index
# rather than a widget
# TODO no indication whether we're currently scrolled up. scroll back
# to bottom after x seconds with no input?
listsize = self.widget.get_item_size(
self.urwid_loop.screen_size, 0, False)
key = self.pane.keypress(listsize, key)
if key:
# `key` gets returned if it wasn't consumed
self.add_log_line(key)
def start(self):
super(DywypiShell, self).start()
#self.hub.network_connected(self.network, self)
def add_log_line(self, line, color='default'):
# TODO generalize this color thing in a way compatible with irc, html, ...
# TODO i super duper want this for logging, showing incoming/outgoing
# messages in the right colors, etc!!
self._print_text((color, line.rstrip()))
def _print_text(self, *encoded_text):
self.pane.body.append(urwid.Text(list(encoded_text)))
self.pane.set_focus(len(self.pane.body) - 1)
# TODO should this just mark dirty??
self.urwid_loop.draw_screen()
def handle_line(self, line):
"""Deal with a line of input."""
try:
self._handle_line(line)
except Exception as e:
log.exception(e)
def _handle_line(self, line):
"""All the good stuff happens here.
Various things happen depending on what the line starts with.
Colon: This is a command; pretend it was sent as a private message.
"""
# Whatever it was, log it
self.pane.body.append(urwid.Text(['>>> ', ('shell-input', line)]))
if line.startswith(':'):
command_string = line[1:]
# TODO rather we didn't need raw_message...
raw_message = ShellMessage(self.me.name, command_string)
event = Message(self, raw_message)
self.event_queue.put_nowait(event)
def _send_message(self, target, message, as_notice=True):
# TODO cool color
self.add_log_line(message)
@asyncio.coroutine
def say(self, target, message):
# TODO target should probably be a peer, eh
if target == self.you.name:
prefix = "bot to you: "
else:
prefix = "bot to {}: ".format(target)
self._print_text(('bot-output-label', prefix), ('bot-output', message))
def format_transition(self, current_style, new_style):
# TODO wait lol shouldn't this be converting to urwid-style tuples
if new_style == Style.default():
# Just use the reset sequence
return '\x1b[0m'
ansi_codes = []
if new_style.fg != current_style.fg:
ansi_codes.append(FOREGROUND_CODES[new_style.fg])
if new_style.bold != current_style.bold:
ansi_codes.append(BOLD_CODES[new_style.bold])
return '\x1b[' + ';'.join(ansi_codes) + 'm'
class GoogleReportState(iot_base.BaseIoT):
"""Report states to Google.
Uses a queue to send messages.
"""
def __init__(self, cloud: Cloud):
"""Initialize Google Report State."""
super().__init__(cloud)
self._connect_lock = asyncio.Lock()
self._to_send = Queue(100)
self._message_sender_task = None
# Local code waiting for a response
self._response_handler: Dict[str, asyncio.Future] = {}
self.register_on_connect(self._async_on_connect)
self.register_on_disconnect(self._async_on_disconnect)
# Register start/stop
cloud.register_on_stop(self.disconnect)
@property
def package_name(self) -> str:
"""Return the package name for logging."""
return __name__
@property
def ws_server_url(self) -> str:
"""Server to connect to."""
# https -> wss, http -> ws
return f"ws{self.cloud.google_actions_report_state_url[4:]}/v1"
async def async_send_message(self, msg):
"""Send a message."""
msgid = uuid.uuid4().hex
# Since connect is async, guard against send_message called twice in parallel.
async with self._connect_lock:
if self.state == iot_base.STATE_DISCONNECTED:
self.cloud.run_task(self.connect())
# Give connect time to start up and change state.
await asyncio.sleep(0)
if self._to_send.full():
discard_msg = self._to_send.get_nowait()
self._response_handler.pop(discard_msg["msgid"]).set_exception(
ErrorResponse(ERR_DISCARD_CODE, ERR_DISCARD_MSG))
fut = self._response_handler[msgid] = asyncio.Future()
self._to_send.put_nowait({"msgid": msgid, "payload": msg})
try:
return await fut
finally:
self._response_handler.pop(msgid, None)
def async_handle_message(self, msg):
"""Handle a message."""
response_handler = self._response_handler.get(msg["msgid"])
if response_handler is not None:
if "error" in msg:
response_handler.set_exception(
ErrorResponse(msg["error"], msg["message"]))
else:
response_handler.set_result(msg.get("payload"))
return
self._logger.warning("Got unhandled message: %s", msg)
async def _async_on_connect(self):
"""On Connect handler."""
self._message_sender_task = self.cloud.run_task(
self._async_message_sender())
async def _async_on_disconnect(self):
"""On disconnect handler."""
self._message_sender_task.cancel()
self._message_sender_task = None
async def _async_message_sender(self):
"""Start sending messages."""
self._logger.debug("Message sender task activated")
try:
while True:
await self.async_send_json_message(await self._to_send.get())
except asyncio.CancelledError:
pass
self._logger.debug("Message sender task shut down")
class HttpProtocol(asyncio.Protocol):
def __init__(self, loop, handler):
self.parser = cparser.HttpRequestParser(
self.on_headers, self.on_body, self.on_error)
self.loop = loop
self.response = Response()
if flavor == 'queue':
def connection_made(self, transport):
self.transport = transport
self.queue = Queue(loop=self.loop)
self.loop.create_task(handle_requests(self.queue, transport))
else:
def connection_made(self, transport):
self.transport = transport
def connection_lost(self, exc):
self.parser.feed_disconnect()
def data_received(self, data):
self.parser.feed(data)
def on_headers(self, request):
return
if flavor == 'block':
def on_body(self, request):
handle_request_block(request, self.transport, self.response)
elif flavor == 'dump':
def on_body(self, request):
handle_dump(request, self.transport, self.response)
elif flavor == 'task':
def on_body(self, request):
self.loop.create_task(handle_request(request, self.transport))
elif flavor == 'queue':
def on_body(self, request):
self.queue.put_nowait(request)
elif flavor == 'inline':
def on_body(self, request):
body = 'Hello inlin!'
status_code = 200
mime_type = 'text/plain'
encoding = 'utf-8'
text = [b'HTTP/1.1 ']
text.extend([str(status_code).encode(), b' OK\r\n'])
text.append(b'Connection: keep-alive\r\n')
text.append(b'Content-Length: ')
text.extend([str(len(body)).encode(), b'\r\n'])
text.extend([
b'Content-Type: ', mime_type.encode(),
b'; encoding=', encoding.encode(), b'\r\n\r\n'])
text.append(body.encode())
self.transport.write(b''.join(text))
elif flavor == 'static':
def on_body(self, request):
self.transport.write(static_response)
def on_error(self, error):
print(error)
class StreamConnection:
def __init__(self, sr, sw, *, loop=None):
if not loop:
loop = asyncio.get_event_loop()
self._loop = loop
self._sr = sr
self._sw = sw
self._msgs = Queue(loop=loop)
self._worker = loop.create_task(self._run())
@asyncio.coroutine
def _run(self):
while self.alive():
try:
data = yield from self._sr.readline()
if data and len(data):
self._msgs.put_nowait(self._convert(data))
except asyncio.CancelledError:
logger.debug("readline from stream reader was cancelled.")
except ConnectionError:
logger.debug("connection error")
break
logger.debug("connection closed")
def _convert(self, data):
return data.strip()
@asyncio.coroutine
def recv(self):
try:
return self._msgs.get_nowait()
except QueueEmpty:
pass
# Wait for a message until the connection is closed
next_message = self._loop.create_task(self._msgs.get())
done, pending = yield from asyncio.wait(
[next_message, self._worker],
loop=self._loop, return_when=asyncio.FIRST_COMPLETED)
if next_message in done:
return next_message.result()
else:
next_message.cancel()
def send(self, data):
if not self.alive():
raise ConnectionError("connection was closed.")
try:
data = data + b'\n'
self._sw.write(data)
except OSError:
raise ConnectionError("can't send data.")
except Exception:
logger.debug("Q___Q")
def alive(self):
return not self._sr.at_eof()
@asyncio.coroutine
def drain():
yield from self._sw.drain()
@asyncio.coroutine
def close(self):
if self.alive():
try:
yield from self._sw.drain()
self._sw.write_eof()
except ConnectionError:
pass
else:
self._sr.feed_eof()
self._sw.close()
self._worker.cancel()
async def radar(radar_Q: Queue = None):
if radar_Q is None:
print("Radar task is terminating. The radar_Q is None. Nothing to do")
return
try:
q = Queue()
while True:
async with async_timeout.timeout(5) as tm:
while True:
try:
incoming_action:INCOMING_ACTION_TYPE = await radar_Q.get()
q.put_nowait(incoming_action)
STATS['number_rx'] += 1
STATS['max_incoming_q_size'] = max(radar_Q.qsize(), STATS['max_incoming_q_size'])
except Exception as x:
break
if not tm.expired:
continue
time_now = datetime.utcnow()
while not q.empty():
item = q.get_nowait()
icao_rec = item.actionMsg
# print(F"DB_WORKER_RXED: {action_types} {icao_rec}")
# await db_process_sbs1_msg2(database, sbs1_msg)
id = item.aircraftId
last_seen = icao_rec['last_seen']
data = icao_rec['current']
delta_sec = (time_now - last_seen).total_seconds()
if delta_sec > 300:
print(F"@Radar: Got an old one: {id} {last_seen} {delta_sec}")
STATS['rx_expired'] += 1
remove_track(id)
del history[id]
continue
history[id] = last_seen
STATS['max_history_size'] = max(STATS['max_history_size'],len(history))
if data['lon'] is not None and data['lat'] is not None:
update_track(id,data['lat'], data['lon'])
STATS['updates'] += 1
update_table_store(id, last_seen, data)
render_table_store()
# endwhile
# cleanup any lingering tracks:
flush()
refresh()
except CancelledError:
print("Cancellling radar task")
except Exception as x:
print(F"Exception {x}")
async def searchWorker(name: str, searchQueue: Queue, queueProgress: tqdm):
while True:
# Get a "work item" out of the queue.
data: QueueObj = await searchQueue.get()
factsLink = data['factsLink']
# url = queryToUrl(query)
print(Fore.MAGENTA + f'{name} got url: {factsLink} to work on' +
Fore.RESET)
start = time.time()
custom_timeout = aiohttp.ClientTimeout(total=60) # type: ignore
# custom_timeout.total = 2*60
async with aiohttp.ClientSession(timeout=custom_timeout) as session:
# nonlocal results
summary_sents = {}
docs_dict = {0: ''}
sents_dict = {0: ''}
isSkipped = True
try:
if (urlparse(factsLink).path.endswith(".pdf") == False):
isSkipped = False
docs_dict, sents_dict = getDocsAndSentsPerUrl(factsLink, 5)
summary_sents = getTfidfSummary(tfidfModel=tfidf,
dictionary=dictionary,
docs_dict=docs_dict,
sents_dict=sents_dict)
except asyncio.TimeoutError as e:
logging.exception(
Fore.RED +
f'Exception raised by worker = {name}; error = {e}' +
Fore.RESET)
except Exception as e: # pylint: disable=broad-except
logging.exception(
Fore.RED +
f'Exception raised by worker = {name}; error = {e}' +
Fore.RESET)
print(
Fore.YELLOW +
f'{name}\'s work for url {factsLink} done; time taken {time.time() - start} seconds'
+ Fore.RESET)
if (isSkipped == True):
print(f'factsLink: \'{factsLink}\' skipped because its pdf'.encode(
encoding='utf-8'))
searchQueue.task_done()
bookRow = dfDeduped[dfDeduped['id'] == data['id']].copy()
bookRow['factsLink'] = factsLink
bookRow['facts'] = [['']]
saveFacts(bookRow)
queueProgress.update(1)
elif (len(sents_dict[0]) > 2):
print(f'factsLink: \'{factsLink}\' fetched with > 2 sentences'.
encode(encoding='utf-8'))
searchQueue.task_done()
bookRow = dfDeduped[dfDeduped['id'] == data['id']].copy()
bookRow['factsLink'] = factsLink
bookRow['facts'] = [[sent for score, sent in summary_sents[0]]]
saveFacts(bookRow)
queueProgress.update(1)
else:
print(
f'factsLink: \'{factsLink}\' unsuccessful, returned < 2 sents, adding to queue again'
.encode(encoding='utf-8'))
searchQueue.task_done()
queueProgress.update(1)
searchQueue.put_nowait(data)
class IRCClient:
"""Higher-level IRC client. Takes care of most of the hard parts of IRC:
incoming server messages are bundled into more intelligible events (see
``dywypi.event``), and commands that expect replies are implemented as
coroutines.
"""
def __init__(self, loop, network):
self.loop = loop
self.network = network
# TODO should this be a param? a property of the network? or, more
# likely, channel-specific and decoded separately and...
self.charset = "utf8"
self.joined_channels = {} # name => Channel
# IRC server features, as reported by ISUPPORT, with defaults taken
# from the RFC.
self.len_nick = 9
self.len_channel = 200
self.len_message = 510
# These lengths don't have limits mentioned in the RFC, so going with
# the smallest known values in the wild
self.len_kick = 80
self.len_topic = 80
self.len_away = 160
self.max_watches = 0
self.max_targets = 1
self.channel_types = set("#&")
self.channel_modes = {} # TODO, haha.
self.channel_prefixes = {} # TODO here too. IRCMode is awkward.
self.network_title = self.network.name
self.features = {}
# Various intermediate state used for waiting for replies and
# aggregating multi-part replies
# TODO hmmm so what happens if state just gets left here forever? do
# we care?
self._pending_names = {}
self._names_futures = {}
self._pending_topics = {}
self._join_futures = {}
self._message_waiters = deque()
self.read_queue = Queue(loop=loop)
def get_channel(self, channel_name):
"""Returns a `Channel` object containing everything the client
definitively knows about the given channel.
Note that if you, say, ask for the topic of a channel you aren't in and
then immediately call `get_channel`, the returned object won't have its
topic populated. State is only tracked persistently for channels the
bot is in; otherwise there's no way to know whether or not it's stale.
"""
if channel_name in self.joined_channels:
return self.joined_channels[channel_name]
else:
return IRCChannel(self, channel_name)
@asyncio.coroutine
def connect(self):
"""Coroutine for connecting to a single server.
Note that this will nonblock until the client is "registered", defined
as the first PING/PONG exchange.
"""
# TODO this is a poor excuse for round-robin :)
server = self.current_server = self.network.servers[0]
# TODO i'm pretty sure the server tells us what our nick is, and we
# should believe that instead
self.nick = self.network.preferred_nick
# TODO: handle disconnection, somehow. probably affects a lot of
# things.
self._reader, self._writer = yield from server.connect(self.loop)
if server.password:
self.send_message("PASS", server.password)
self.send_message("NICK", self.nick)
self.send_message("USER", "dywypi", "-", "-", "dywypi Python IRC bot")
# Start the reader loop, or we can't respond to anything
self._read_loop_task = asyncio.Task(self._start_read_loop())
asyncio.async(self._read_loop_task, loop=self.loop)
@asyncio.coroutine
def disconnect(self):
# Quit
self.send_message("QUIT", "Seeya!")
# Flush the write buffer
yield from self._writer.drain()
self._writer.close()
# Stop reading events
self._read_loop_task.cancel()
# This looks a little funny since this task is already running, but we
# want to block until it's actually done, which might require dipping
# back into the event loop
yield from self._read_loop_task
# Read until the connection closes
while not self._reader.at_eof():
yield from self._reader.readline()
@asyncio.coroutine
def _start_read_loop(self):
"""Internal coroutine that just keeps reading from the server in a
loop. Called once after a connect and should never be called again
after that.
"""
# TODO this is currently just to keep the message queue going, but
# eventually it should turn them into events and stuff them in an event
# queue
while not self._reader.at_eof():
try:
yield from self._read_message()
except CancelledError:
return
except Exception:
log.exception("Smothering exception in IRC read loop")
@asyncio.coroutine
def gather_messages(self, *start, finish):
fut = asyncio.Future()
messages = {}
for command in start:
messages[command] = False
for command in finish:
messages[command] = True
collected = []
self._message_waiters.append((fut, messages, collected))
yield from fut
return collected
def _possibly_gather_message(self, message):
if not self._message_waiters:
return
# TODO there is a general ongoing problem here with matching up
# responses. ESPECIALLY when error codes are possible. something here
# is gonna have to get a bit fancier.
fut, waiting_on, collected = self._message_waiters[0]
# TODO is it possible for even a PING to appear in the middle of
# some other response?
# TODO this is still susceptible to weirdness when there's, say, a
# queued error response to a PRIVMSG on its way back; it'll look
# like the call we just made failed, and all the real responses
# will be dropped. can we assume some set of error replies ONLY
# happen in response to sending a message of some kind, maybe?
# TODO for that matter, where does the error response to a PRIVMSG
# even go? the whole problem is that we can't know for sure when
# it succeeded, unless we put a timeout on every call to say()
finish = False
if message.command in waiting_on:
finish = waiting_on[message.command]
elif message.is_error:
# Always consider an error as finishing
# TODO but we might have gotten this error in response to something
# else we did before this message... :S
if message.command in {"ERR_CANNOTSENDTOCHAN"}:
# Looks like a PRIVMSG error or similar, so probably not a
# response to this particular message.
return
finish = True
elif not collected:
# Got a regular response we weren't expecting, AND this future
# hasn't started collecting yet -- the response probably just
# hasn't started coming back yet, so don't do anything yet.
return
# If we get here, we expected this response, and should keep
# feeding into this future.
collected.append(message)
if finish:
# Done, one way or another
self._message_waiters.popleft()
if message.is_error:
fut.set_exception(IRCError(message))
else:
fut.set_result(collected)
@asyncio.coroutine
def _read_message(self):
"""Internal dispatcher for messages received from the server."""
line = yield from self._reader.readline()
assert line.endswith(b"\r\n")
line = line[:-2]
# TODO valerr, unicodeerr
message = IRCMessage.parse(line.decode(self.charset))
log.debug("recv: %r", message)
# TODO unclear whether this should go before or after _handle_foo
self._possibly_gather_message(message)
# Boy do I ever hate this pattern but it's slightly more maintainable
# than a 500-line if tree.
handler = getattr(self, "_handle_" + message.command, None)
event = None
if handler:
event = handler(message)
self.read_queue.put_nowait((message, event))
def _handle_PING(self, message):
# PONG
self.send_message("PONG", message.args[-1])
def _handle_RPL_WELCOME(self, message):
# Initial registration: do autojoins, and any other onconnect work
for channel_name in self.network.autojoins:
asyncio.async(self.join(channel_name), loop=self.loop)
def _handle_RPL_ISUPPORT(self, message):
me, *features, human_text = message.args
for feature_string in features:
feature, _, value = feature_string.partition("=")
if value is None:
value = True
self.features[feature] = value
if feature == "NICKLEN":
self.len_nick = int(value)
elif feature == "CHANNELLEN":
self.len_channel = int(value)
elif feature == "KICKLEN":
self.len_kick = int(value)
elif feature == "TOPICLEN":
self.len_topic = int(value)
elif feature == "AWAYLEN":
self.len_away = int(value)
elif feature == "WATCH":
self.max_watches = int(value)
elif feature == "CHANTYPES":
self.channel_types = set(value)
elif feature == "PREFIX":
# List of channel user modes, in relative priority order, in
# the format (ov)@+
assert value[0] == "("
letters, symbols = value[1:].split(")")
assert len(letters) == len(symbols)
self.channel_prefixes.clear()
for letter, symbol in zip(letters, symbols):
mode = IRCMode(letter, prefix=symbol)
self.channel_modes[letter] = mode
self.channel_prefixes[symbol] = mode
elif feature == "MAXTARGETS":
self.max_targets = int(value)
elif feature == "CHANMODES":
# Four groups delimited by lists: list-style (+b), arg required
# (+k), arg required only to set (+l), argless
lists, args, argsets, argless = value.split(",")
for letter in lists:
self.channel_modes[letter] = IRCMode(letter, multi=True)
for letter in args:
self.channel_modes[letter] = IRCMode(letter, arg_on_set=True, arg_on_remove=True)
for letter in argsets:
self.channel_modes[letter] = IRCMode(letter, arg_on_set=True)
for letter in argless:
self.channel_modes[letter] = IRCMode(letter)
elif feature == "NETWORK":
self.network_title = value
def _handle_JOIN(self, message):
channel_name, = message.args
joiner = Peer.from_prefix(message.prefix)
# TODO should there be a self.me? how...
if joiner.name == self.nick:
# We just joined a channel
# assert channel_name not in self.joined_channels
# TODO key? do we care?
# TODO what about channel configuration and anon non-joined
# channels? how do these all relate...
channel = IRCChannel(self, channel_name)
self.joined_channels[channel.name] = channel
else:
# Someone else just joined the channel
self.joined_channels[channel_name].add_user(joiner)
def _handle_RPL_TOPIC(self, message):
# Topic. Sent when joining or when requesting the topic.
# TODO this doesn't handle the "requesting" part
# TODO what if me != me?
me, channel_name, topic_text = message.args
self._pending_topics[channel_name] = IRCTopic(topic_text)
def _handle_RPL_TOPICWHOTIME(self, message):
# Topic author (NONSTANDARD). Sent after RPL_TOPIC.
# Unfortunately, there's no way to know whether to expect this.
# TODO this doesn't handle the "requesting" part
# TODO what if me != me?
me, channel_name, author, timestamp = message.args
topic = self._pending_topics.setdefault(channel_name, IRCTopic(""))
topic.author = Peer.from_prefix(author)
topic.timestamp = datetime.utcfromtimestamp(int(timestamp))
def _handle_RPL_NAMREPLY(self, message):
# Names response. Sent when joining or when requesting a names
# list. Must be ended with a RPL_ENDOFNAMES.
me, useless_equals_sign, channel_name, *raw_names = message.args
# List of names is actually optional (?!)
if raw_names:
raw_names = raw_names[0]
else:
raw_names = ""
names = raw_names.strip(" ").split(" ")
namelist = self._pending_names.setdefault(channel_name, [])
# TODO modes? should those be stripped off here?
# TODO for that matter should these become peers here?
namelist.extend(names)
def _handle_RPL_ENDOFNAMES(self, message):
# End of names list. Sent at the very end of a join or the very
# end of a NAMES request.
me, channel_name, info = message.args
namelist = self._pending_names.pop(channel_name, [])
if channel_name in self._names_futures:
# TODO we should probably not ever have a names future AND a
# pending join at the same time. or, does it matter?
self._names_futures[channel_name].set_result(namelist)
del self._names_futures[channel_name]
if channel_name in self.joined_channels:
# Join synchronized!
channel = self.joined_channels[channel_name]
channel.sync = True
channel.topic = self._pending_topics.pop(channel_name, None)
for name in namelist:
modes = set()
# TODO use features!
while name and name[0] in "+%@&~":
modes.add(name[0])
name = name[1:]
# TODO haha no this is so bad.
# TODO the bot should, obviously, keep a record of all
# known users as well. alas, mutable everything.
peer = Peer(name, None, None)
channel.add_user(peer, modes)
if channel_name in self._join_futures:
# Update the Future
self._join_futures[channel_name].set_result(channel)
del self._join_futures[channel_name]
def _handle_PRIVMSG(self, message):
# PRIVMSG target :text
target_name, text = message.args
source = Peer.from_prefix(message.prefix)
if target_name[0] in self.channel_types:
target = self.get_channel(target_name)
cls = PublicMessage
else:
# TODO this is /us/, so, surely ought to be known
target = Peer(target_name, None, None)
cls = PrivateMessage
return cls(source, target, text, client=self, raw=message)
@asyncio.coroutine
def read_event(self):
"""Produce a single IRC event.
This client does not do any kind of multiplexing or event handler
notification; that's left to a higher level.
"""
message, event = yield from self.read_queue.get()
return event
# Implementations of particular commands
# TODO should these be part of the general client interface, or should
# there be a separate thing that smooths out the details?
@asyncio.coroutine
def whois(self, target):
"""Coroutine that queries for information about a target."""
self.send_message("WHOIS", target)
messages = yield from self.gather_messages(
"RPL_WHOISUSER",
"RPL_WHOISSERVER",
"RPL_WHOISOPERATOR",
"RPL_WHOISIDLE",
"RPL_WHOISCHANNELS",
"RPL_WHOISVIRT",
"RPL_WHOIS_HIDDEN",
"RPL_WHOISSPECIAL",
"RPL_WHOISSECURE",
"RPL_WHOISSTAFF",
"RPL_WHOISLANGUAGE",
finish=["RPL_ENDOFWHOIS", "ERR_NOSUCHSERVER", "ERR_NONICKNAMEGIVEN", "ERR_NOSUCHNICK"],
)
# nb: The first two args for all the responses are our nick and the
# target's nick.
# TODO apparently you can whois multiple nicks at a time
for message in messages:
if message.command == "RPL_WHOISUSER":
ident = message.args[2]
hostname = message.args[3]
# args[4] is a literal *
realname = message.args[5]
elif message.command == "RPL_WHOISIDLE":
# Idle time. Some servers (at least, inspircd) also have
# signon time as unixtime.
idle = timedelta(seconds=int(message.args[2]))
elif message.command == "RPL_WHOISCHANNELS":
# TODO split and parse out the usermodes
# TODO don't some servers have an extension with multiple modes
# here
channels = message.args[2]
elif message.command == "RPL_WHOISSERVER":
server = message.args[2]
server_desc = message.args[3]
return messages
@asyncio.coroutine
def say(self, target, message):
"""Coroutine that sends a message to a target, which may be either a
`Channel` or a `Peer`.
"""
self.send_message("PRIVMSG", target, message)
@asyncio.coroutine
def join(self, channel_name, key=None):
"""Coroutine that joins a channel, and nonblocks until the join is
"synchronized" (defined as receiving the nick list).
"""
if channel_name in self._join_futures:
return self._join_futures[channel_name]
# TODO multiple? error on commas?
if key is None:
self.send_message("JOIN", channel_name)
else:
self.send_message("JOIN", channel_name, key)
# Clear out any lingering names list
self._pending_names[channel_name] = []
# Return a Future, to be populated by the message loop
fut = self._join_futures[channel_name] = asyncio.Future()
return fut
@asyncio.coroutine
def names(self, channel_name):
"""Coroutine that returns a list of names in a channel."""
# TODO there's some ISUPPORT extension that lists /all/ channel modes
# on each name that comes back... support that?
self.send_message("NAMES", channel_name)
# No need to do the same thing twice
if channel_name in self._names_futures:
return self._names_futures[channel_name]
# Clear out any lingering names list
self._pending_names[channel_name] = []
# Return a Future, to be populated by the message loop
fut = self._names_futures[channel_name] = asyncio.Future()
return fut
def set_topic(self, channel, topic):
"""Sets the channel topic."""
self.send_message("TOPIC", channel, topic)
# TODO unclear whether this stuff should be separate or what; it's less
# about the protocol and more about the dywypi interface
def send_message(self, command, *args):
message = IRCMessage(command, *args)
log.debug("sent: %r", message)
self._writer.write(message.render().encode(self.charset) + b"\r\n")
def format_transition(self, current_style, new_style):
if new_style == Style.default():
# Reset code, ^O
return "\x0f"
if new_style.fg != current_style.fg and new_style.fg is Color.default:
# IRC has no "reset to default" code. mIRC claims color 99 is for
# this, but it lies, at least in irssi. So we must reset and
# reapply everything.
ret = "\x0f"
if new_style.bold is Bold.on:
ret += "\x02"
return ret
ret = ""
if new_style.fg != current_style.fg:
ret += FOREGROUND_CODES[new_style.fg]
if new_style.bold != current_style.bold:
# There's no on/off for bold, just a toggle
ret += "\x02"
return ret
class IRCClient:
"""Higher-level IRC client. Takes care of most of the hard parts of IRC:
incoming server messages are bundled into more intelligible events (see
``dywypi.event``), and commands that expect replies are implemented as
coroutines.
"""
def __init__(self, loop, network):
self.loop = loop
self.network = network
# TODO should this be a param? a property of the network? or, more
# likely, channel-specific and decoded separately and...
self.charset = 'utf8'
self.joined_channels = {} # name => Channel
# IRC server features, as reported by ISUPPORT, with defaults taken
# from the RFC.
self.len_nick = 9
self.len_channel = 200
self.len_message = 510
# These lengths don't have limits mentioned in the RFC, so going with
# the smallest known values in the wild
self.len_kick = 80
self.len_topic = 80
self.len_away = 160
self.max_watches = 0
self.max_targets = 1
self.channel_types = set('#&')
self.channel_modes = {} # TODO, haha.
self.channel_prefixes = {} # TODO here too. IRCMode is awkward.
self.network_title = self.network.name
self.features = {}
# Various intermediate state used for waiting for replies and
# aggregating multi-part replies
# TODO hmmm so what happens if state just gets left here forever? do
# we care?
self._pending_names = {}
self._names_futures = {}
self._pending_topics = {}
self._join_futures = {}
self._message_waiters = OrderedDict()
self.read_queue = Queue(loop=loop)
def get_channel(self, channel_name):
"""Returns a `Channel` object containing everything the client
definitively knows about the given channel.
Note that if you, say, ask for the topic of a channel you aren't in and
then immediately call `get_channel`, the returned object won't have its
topic populated. State is only tracked persistently for channels the
bot is in; otherwise there's no way to know whether or not it's stale.
"""
if channel_name in self.joined_channels:
return self.joined_channels[channel_name]
else:
return IRCChannel(self, channel_name)
@asyncio.coroutine
def connect(self):
"""Coroutine for connecting to a single server.
Note that this will nonblock until the client is "registered", defined
as the first PING/PONG exchange.
"""
# TODO this is a poor excuse for round-robin :)
server = self.current_server = self.network.servers[0]
# TODO i'm pretty sure the server tells us what our nick is, and we
# should believe that instead
self.nick = self.network.preferred_nick
# TODO: handle disconnection, somehow. probably affects a lot of
# things.
self._reader, self._writer = yield from server.connect(self.loop)
log.debug('connected!')
if server.password:
self.send_message('PASS', server.password)
self.send_message('NICK', self.nick)
self.send_message('USER', 'dywypi', '-', '-', 'dywypi Python IRC bot')
# Start the reader loop, or we can't respond to anything
self._read_loop_task = asyncio.Task(self._start_read_loop())
asyncio.async(self._read_loop_task, loop=self.loop)
@asyncio.coroutine
def disconnect(self):
# Quit
self.send_message('QUIT', 'Seeya!')
# Flush the write buffer
yield from self._writer.drain()
self._writer.close()
# Stop reading events
self._read_loop_task.cancel()
# This looks a little funny since this task is already running, but we
# want to block until it's actually done, which might require dipping
# back into the event loop
yield from self._read_loop_task
# Read until the connection closes
while not self._reader.at_eof():
yield from self._reader.readline()
@asyncio.coroutine
def _start_read_loop(self):
"""Internal coroutine that just keeps reading from the server in a
loop. Called once after a connect and should never be called again
after that.
"""
# TODO this is currently just to keep the message queue going, but
# eventually it should turn them into events and stuff them in an event
# queue
while not self._reader.at_eof():
try:
yield from self._read_message()
except CancelledError:
return
except Exception:
log.exception("Smothering exception in IRC read loop")
@asyncio.coroutine
def gather_messages(self, *middle, end, errors=()):
fut = asyncio.Future()
messages = {}
for command in middle:
messages[command] = 'middle'
for command in end:
messages[command] = 'end'
for command in errors:
messages[command] = 'error'
collected = []
self._message_waiters[fut] = (messages, collected)
yield from fut
return collected
@asyncio.coroutine
def _read_message(self):
"""Internal dispatcher for messages received from the server."""
line = yield from self._reader.readline()
assert line.endswith(b'\r\n')
line = line[:-2]
# TODO valerr, unicodeerr
message = IRCMessage.parse(line.decode(self.charset))
log.debug("recv: %r", message)
# TODO there is a general ongoing problem here with matching up
# responses. ESPECIALLY when error codes are possible. something here
# is gonna have to get a bit fancier.
for fut, (waiting_on, collected) in self._message_waiters.items():
# TODO this needs to handle error codes too, or the future will
# linger forever! potential problem: if the server is lagging
# behind us, an error code might actually map to a privmsg we tried
# to send (which has no success response) and we'll get all f****d
# up. i don't know if there's any way to solve this.
# TODO hey stupid question: after we've seen ANY of the waited-on
# messages, should we pipe all subsequent messages into that future
# until we see the one that's supposed to end it? something like
# a forced JOIN could screw up a join attempt, for example, but if
# we're getting RPL_TOPIC when we didn't actually ask for the
# topic, THEN we know we're definitely in the join sequence.
# TODO also given normal irc response flow, i'm pretty sure we
# should only ever need to check the first pending future. there's
# no way we should need to skip around.
# TODO maybe give these a timeout so a bad one doesn't f**k us up
# forever
if message.command in waiting_on:
collected.append(message)
if waiting_on[message.command] == 'end':
fut.set_result(collected)
del self._message_waiters[fut]
elif waiting_on[message.command] == 'error':
fut.set_exception(IRCError(message))
del self._message_waiters[fut]
break
# Boy do I ever hate this pattern but it's slightly more maintainable
# than a 500-line if tree.
handler = getattr(self, '_handle_' + message.command, None)
event = None
if handler:
event = handler(message)
self.read_queue.put_nowait((message, event))
def _handle_PING(self, message):
# PONG
self.send_message('PONG', message.args[-1])
def _handle_RPL_WELCOME(self, message):
# Initial registration: do autojoins, and any other onconnect work
self.network.hostname = message.args[1].rsplit(sep='@')[-1]
for channel_name in self.network.autojoins:
asyncio.async(self.join(channel_name), loop=self.loop)
def _handle_RPL_ISUPPORT(self, message):
me, *features, human_text = message.args
for feature_string in features:
feature, _, value = feature_string.partition('=')
if value is None:
value = True
self.features[feature] = value
if feature == 'NICKLEN':
self.len_nick = int(value)
elif feature == 'CHANNELLEN':
self.len_channel = int(value)
elif feature == 'KICKLEN':
self.len_kick = int(value)
elif feature == 'TOPICLEN':
self.len_topic = int(value)
elif feature == 'AWAYLEN':
self.len_away = int(value)
elif feature == 'WATCH':
self.max_watches = int(value)
elif feature == 'CHANTYPES':
self.channel_types = set(value)
elif feature == 'PREFIX':
# List of channel user modes, in relative priority order, in
# the format (ov)@+
assert value[0] == '('
letters, symbols = value[1:].split(')')
assert len(letters) == len(symbols)
self.channel_prefixes.clear()
for letter, symbol in zip(letters, symbols):
mode = IRCMode(letter, prefix=symbol)
self.channel_modes[letter] = mode
self.channel_prefixes[symbol] = mode
elif feature == 'MAXTARGETS':
self.max_targets = int(value)
elif feature == 'CHANMODES':
# Four groups delimited by lists: list-style (+b), arg required
# (+k), arg required only to set (+l), argless
lists, args, argsets, argless = value.split(',')
for letter in lists:
self.channel_modes[letter] = IRCMode(
letter, multi=True)
for letter in args:
self.channel_modes[letter] = IRCMode(
letter, arg_on_set=True, arg_on_remove=True)
for letter in argsets:
self.channel_modes[letter] = IRCMode(
letter, arg_on_set=True)
for letter in argless:
self.channel_modes[letter] = IRCMode(letter)
elif feature == 'NETWORK':
self.network_title = value
def _handle_JOIN(self, message):
channel_name, = message.args
joiner = Peer.from_prefix(message.prefix)
# TODO should there be a self.me? how...
if joiner.name == self.nick:
# We just joined a channel
#assert channel_name not in self.joined_channels
# TODO key? do we care?
# TODO what about channel configuration and anon non-joined
# channels? how do these all relate...
channel = IRCChannel(self, channel_name)
self.joined_channels[channel.name] = channel
else:
# Someone else just joined the channel
self.joined_channels[channel_name].add_user(joiner)
def _handle_RPL_TOPIC(self, message):
# Topic. Sent when joining or when requesting the topic.
# TODO this doesn't handle the "requesting" part
# TODO what if me != me?
me, channel_name, topic_text = message.args
self._pending_topics[channel_name] = IRCTopic(topic_text)
def _handle_RPL_TOPICWHOTIME(self, message):
# Topic author (NONSTANDARD). Sent after RPL_TOPIC.
# Unfortunately, there's no way to know whether to expect this.
# TODO this doesn't handle the "requesting" part
# TODO what if me != me?
me, channel_name, author, timestamp = message.args
topic = self._pending_topics.setdefault(channel_name, IRCTopic(''))
topic.author = Peer.from_prefix(author)
topic.timestamp = datetime.utcfromtimestamp(int(timestamp))
def _handle_RPL_NAMREPLY(self, message):
# Names response. Sent when joining or when requesting a names
# list. Must be ended with a RPL_ENDOFNAMES.
me, useless_equals_sign, channel_name, *raw_names = message.args
# List of names is actually optional (?!)
if raw_names:
raw_names = raw_names[0]
else:
raw_names = ''
names = raw_names.strip(' ').split(' ')
namelist = self._pending_names.setdefault(channel_name, [])
# TODO modes? should those be stripped off here?
# TODO for that matter should these become peers here?
namelist.extend(names)
def _handle_RPL_ENDOFNAMES(self, message):
# End of names list. Sent at the very end of a join or the very
# end of a NAMES request.
me, channel_name, info = message.args
namelist = self._pending_names.pop(channel_name, [])
if channel_name in self._names_futures:
# TODO we should probably not ever have a names future AND a
# pending join at the same time. or, does it matter?
self._names_futures[channel_name].set_result(namelist)
del self._names_futures[channel_name]
if channel_name in self.joined_channels:
# Join synchronized!
channel = self.joined_channels[channel_name]
channel.sync = True
channel.topic = self._pending_topics.pop(channel_name, None)
for name in namelist:
modes = set()
# TODO use features!
while name and name[0] in '+%@&~':
modes.add(name[0])
name = name[1:]
# TODO haha no this is so bad.
# TODO the bot should, obviously, keep a record of all
# known users as well. alas, mutable everything.
peer = Peer(name, None, None)
channel.add_user(peer, modes)
if channel_name in self._join_futures:
# Update the Future
self._join_futures[channel_name].set_result(channel)
del self._join_futures[channel_name]
def _handle_PRIVMSG(self, message):
# PRIVMSG target :text
target_name, text = message.args
source = Peer.from_prefix(message.prefix)
if target_name[0] in self.channel_types:
target = self.get_channel(target_name)
cls = PublicMessage
else:
# TODO this is /us/, so, surely ought to be known
target = Peer(target_name, None, None)
cls = PrivateMessage
return cls(source, target, text, client=self, raw=message)
@asyncio.coroutine
def read_event(self):
"""Produce a single IRC event.
This client does not do any kind of multiplexing or event handler
notification; that's left to a higher level.
"""
message, event = yield from self.read_queue.get()
return event
# Implementations of particular commands
# TODO should these be part of the general client interface, or should
# there be a separate thing that smooths out the details?
@asyncio.coroutine
def whois(self, target):
"""Coroutine that queries for information about a target."""
self.send_message('WHOIS', target)
messages = yield from self.gather_messages(
'RPL_WHOISUSER',
'RPL_WHOISSERVER',
'RPL_WHOISOPERATOR',
'RPL_WHOISIDLE',
'RPL_WHOISCHANNELS',
'RPL_WHOISVIRT',
'RPL_WHOIS_HIDDEN',
'RPL_WHOISSPECIAL',
'RPL_WHOISSECURE',
'RPL_WHOISSTAFF',
'RPL_WHOISLANGUAGE',
end=[
'RPL_ENDOFWHOIS',
],
errors=[
'ERR_NOSUCHSERVER',
'ERR_NONICKNAMEGIVEN',
'ERR_NOSUCHNICK',
],
)
# nb: The first two args for all the responses are our nick and the
# target's nick.
# TODO apparently you can whois multiple nicks at a time
for message in messages:
if message.command == 'RPL_WHOISUSER':
ident, hostname
ident = message.args[2]
hostname = message.args[3]
# args[4] is a literal *
realname = message.args[5]
elif message.command == 'RPL_WHOISIDLE':
# Idle time. Some servers (at least, inspircd) also have
# signon time as unixtime.
idle = timedelta(seconds=int(message.args[2]))
elif message.command == 'RPL_WHOISCHANNELS':
# TODO split and parse out the usermodes
# TODO don't some servers have an extension with multiple modes
# here
channels = message.args[2]
elif message.command == 'RPL_WHOISSERVER':
server = message.args[2]
server_desc = message.args[3]
return messages
@asyncio.coroutine
def say(self, message, target, notice=False):
"""Coroutine that sends a message to a target, which may be either a
`Channel` or a `Peer`.
"""
command = 'NOTICE' if notice else 'PRIVMSG'
self.send_message(command, target, message)
@asyncio.coroutine
def join(self, channel_name, key=None):
"""Coroutine that joins a channel, and nonblocks until the join is
"synchronized" (defined as receiving the nick list).
"""
if channel_name in self._join_futures:
return self._join_futures[channel_name]
# TODO multiple? error on commas?
if key is None:
self.send_message('JOIN', channel_name)
else:
self.send_message('JOIN', channel_name, key)
# Clear out any lingering names list
self._pending_names[channel_name] = []
# Return a Future, to be populated by the message loop
fut = self._join_futures[channel_name] = asyncio.Future()
return fut
@asyncio.coroutine
def names(self, channel_name):
"""Coroutine that returns a list of names in a channel."""
self.send_message('NAMES', channel_name)
# No need to do the same thing twice
if channel_name in self._names_futures:
return self._names_futures[channel_name]
# Clear out any lingering names list
self._pending_names[channel_name] = []
# Return a Future, to be populated by the message loop
fut = self._names_futures[channel_name] = asyncio.Future()
return fut
def set_topic(self, channel, topic):
"""Sets the channel topic."""
self.send_message('TOPIC', channel, topic)
# TODO unclear whether this stuff should be separate or what; it's less
# about the protocol and more about the dywypi interface
def send_message(self, command, *args):
message = IRCMessage(command, *args)
log.debug("sent: %r", message)
self._writer.write(message.render().encode(self.charset) + b'\r\n')
def format_transition(self, current_style, new_style):
if new_style == Style.default():
# Reset code, ^O
return '\x0f'
if new_style.fg != current_style.fg and new_style.fg is Color.default:
# IRC has no "reset to default" code. mIRC claims color 99 is for
# this, but it lies, at least in irssi. So we must reset and
# reapply everything.
ret = '\x0f'
if new_style.bold is Bold.on:
ret += '\x02'
return ret
ret = ''
if new_style.fg != current_style.fg:
ret += FOREGROUND_CODES[new_style.fg]
if new_style.bold != current_style.bold:
# There's no on/off for bold, just a toggle
ret += '\x02'
return ret
class Ant(abc.ABC):
response_pipelines: typing.List[Pipeline] = []
request_pipelines: typing.List[Pipeline] = []
item_pipelines: typing.List[Pipeline] = []
request_cls = Request
response_cls = Response
request_timeout = DEFAULT_TIMEOUT.total
request_retries = 3
request_retry_delay = 5
request_proxies: typing.List[typing.Union[str, URL]] = []
request_max_redirects = 10
request_allow_redirects = True
response_in_stream = False
connection_limit = 100 # see "TCPConnector" in "aiohttp"
connection_limit_per_host = 0
concurrent_limit = 100
def __init__(self,
loop: typing.Optional[asyncio.AbstractEventLoop] = None):
self.loop = loop if loop is not None else asyncio.get_event_loop()
self.logger = logging.getLogger(self.__class__.__name__)
self.session: aiohttp.ClientSession = ClientSession(
response_class=self.response_cls,
connector=aiohttp.TCPConnector(
limit=self.connection_limit,
enable_cleanup_closed=True,
limit_per_host=self.connection_limit_per_host))
# coroutine`s concurrency support
self._queue = Queue(loop=self.loop)
self._done_queue = Queue(loop=self.loop)
self._running_count = 0
self._is_closed = False
# report var
self._reports: typing.DefaultDict[str, typing.List[
int, int]] = defaultdict(lambda: [0, 0])
self._drop_reports: typing.DefaultDict[str, typing.List[
int, int]] = defaultdict(lambda: [0, 0])
self._start_time = time.time()
self._last_time = self._start_time
self._report_slot = 60 # report once after one minute by default
@property
def name(self):
return self.__class__.__name__
@property
def is_running(self) -> bool:
return self._running_count > 0
async def request(
self,
url: typing.Union[str, URL],
method: str = aiohttp.hdrs.METH_GET,
params: typing.Optional[dict] = None,
headers: typing.Optional[dict] = None,
cookies: typing.Optional[dict] = None,
data: typing.Optional[typing.Union[typing.AnyStr, typing.Dict,
typing.IO]] = None,
proxy: typing.Optional[typing.Union[str, URL]] = None,
timeout: typing.Optional[typing.Union[int, float]] = None,
retries: typing.Optional[int] = None,
response_in_stream: typing.Optional[bool] = None) -> Response:
if not isinstance(url, URL):
url = URL(url)
if proxy and not isinstance(proxy, URL):
proxy = URL(proxy)
elif proxy is None:
proxy = self.get_proxy()
if timeout is None:
timeout = self.request_timeout
if retries is None:
retries = self.request_retries
if response_in_stream is None:
response_in_stream = self.response_in_stream
req = self.request_cls(method,
url,
timeout=timeout,
params=params,
headers=headers,
cookies=cookies,
data=data,
proxy=proxy,
response_in_stream=response_in_stream)
req = await self._handle_thing_with_pipelines(req,
self.request_pipelines)
self.report(req)
if retries > 0:
res = await self.make_retry_decorator(
retries, self.request_retry_delay)(self._request)(req)
else:
res = await self._request(req)
res = await self._handle_thing_with_pipelines(res,
self.response_pipelines)
self.report(res)
return res
async def collect(self, item: Item) -> None:
self.logger.debug('Collect item: ' + str(item))
await self._handle_thing_with_pipelines(item, self.item_pipelines)
self.report(item)
async def open(self) -> None:
self.logger.info('Opening')
for pipeline in itertools.chain(self.item_pipelines,
self.response_pipelines,
self.request_pipelines):
obj = pipeline.on_spider_open()
if asyncio.iscoroutine(obj):
await obj
async def close(self) -> None:
await self.wait_scheduled_coroutines()
for pipeline in itertools.chain(self.item_pipelines,
self.response_pipelines,
self.request_pipelines):
obj = pipeline.on_spider_close()
if asyncio.iscoroutine(obj):
await obj
await self.session.close()
self._is_closed = True
self.logger.info('Closed')
@abc.abstractmethod
async def run(self) -> None:
"""App custom entrance"""
async def main(self) -> None:
try:
await self.open()
await self.run()
except Exception as e:
self.logger.exception('Run ant with ' + e.__class__.__name__)
try:
await self.close()
except Exception as e:
self.logger.exception('Close ant with ' + e.__class__.__name__)
# total report
for name, counts in self._reports.items():
self.logger.info('Get {:d} {:s} in total'.format(counts[1], name))
for name, counts in self._drop_reports.items():
self.logger.info('Drop {:d} {:s} in total'.format(counts[1], name))
self.logger.info('Run {:s} in {:f} seconds'.format(
self.__class__.__name__,
time.time() - self._start_time))
@staticmethod
def make_retry_decorator(
retries: int, delay: float
) -> typing.Callable[[typing.Callable], typing.Callable]:
return retry(wait=wait_fixed(delay),
retry=(retry_if_result(lambda res: res.status >= 500)
| retry_if_exception_type(
exception_types=aiohttp.ClientError)),
stop=stop_after_attempt(retries + 1))
def get_proxy(self) -> typing.Optional[URL]:
"""Chose a proxy, default by random"""
try:
return URL(random.choice(self.request_proxies))
except IndexError:
return None
def schedule_coroutine(self, coroutine: typing.Coroutine) -> None:
"""Like "asyncio.ensure_future", it schedule coroutine in event loop
and return immediately.
Call "self.wait_scheduled_coroutines" make sure all coroutine has been
done.
"""
def _done_callback(f):
self._running_count -= 1
self._done_queue.put_nowait(f)
try:
if (self.concurrent_limit == -1
or self._running_count < self.concurrent_limit):
next_coroutine = self._queue.get_nowait()
self._running_count += 1
asyncio.ensure_future(
next_coroutine,
loop=self.loop).add_done_callback(_done_callback)
except QueueEmpty:
pass
if self._is_closed:
self.logger.warning('This pool has be closed!')
return
if (self.concurrent_limit == -1
or self._running_count < self.concurrent_limit):
self._running_count += 1
asyncio.ensure_future(
coroutine, loop=self.loop).add_done_callback(_done_callback)
else:
self._queue.put_nowait(coroutine)
def schedule_coroutines(
self, coroutines: typing.Iterable[typing.Coroutine]) -> None:
"""A short way to schedule many coroutines.
"""
for coroutine in coroutines:
self.schedule_coroutine(coroutine)
async def wait_scheduled_coroutines(self):
"""Wait scheduled coroutines to be done, can be called many times.
"""
while self._running_count > 0 or self._done_queue.qsize() > 0:
await self._done_queue.get()
def as_completed(
self,
coroutines: typing.Iterable[typing.Coroutine],
limit: typing.Optional[int] = None
) -> typing.Generator[typing.Coroutine, None, None]:
"""Like "asyncio.as_completed",
run and iter coroutines out of the pool.
:param limit: set to "self.concurrent_limit" by default,
this "limit" is not shared with pool`s limit
"""
limit = self.concurrent_limit if limit is None else limit
coroutines = iter(coroutines)
queue = Queue(loop=self.loop)
todo = []
def _done_callback(f):
queue.put_nowait(f)
todo.remove(f)
try:
nf = asyncio.ensure_future(next(coroutines))
nf.add_done_callback(_done_callback)
todo.append(nf)
except StopIteration:
pass
async def _wait_for_one():
f = await queue.get()
return f.result()
if limit <= 0:
fs = {
asyncio.ensure_future(cor, loop=self.loop)
for cor in coroutines
}
else:
fs = {
asyncio.ensure_future(cor, loop=self.loop)
for cor in islice(coroutines, 0, limit)
}
for f in fs:
f.add_done_callback(_done_callback)
todo.append(f)
while len(todo) > 0 or queue.qsize() > 0:
yield _wait_for_one()
async def as_completed_with_async(
self,
coroutines: typing.Iterable[typing.Coroutine],
limit: typing.Optional[int] = None,
raise_exception: bool = True,
) -> typing.AsyncGenerator[typing.Any, None]:
"""as_completed`s async version, can catch and log exception inside.
"""
for coro in self.as_completed(coroutines, limit=limit):
try:
yield await coro
except Exception as e:
if raise_exception:
raise e
else:
self.logger.exception('Get exception {:s} in '
'"as_completed_with_async"'.format(
str(e)))
def report(self, thing: Things, dropped: bool = False) -> None:
now_time = time.time()
if now_time - self._last_time > self._report_slot:
self._last_time = now_time
for name, counts in self._reports.items():
count = counts[1] - counts[0]
counts[0] = counts[1]
self.logger.info(
'Get {:d} {:s} in total with {:d}/{:d}s rate'.format(
counts[1], name, count, self._report_slot))
for name, counts in self._drop_reports.items():
count = counts[1] - counts[0]
counts[0] = counts[1]
self.logger.info(
'Drop {:d} {:s} in total with {:d}/{:d} rate'.format(
counts[1], name, count, self._report_slot))
report_type = thing.__class__.__name__
if dropped:
reports = self._drop_reports
else:
reports = self._reports
counts = reports[report_type]
counts[1] += 1
async def _handle_thing_with_pipelines(
self, thing: Things, pipelines: typing.List[Pipeline]) -> Things:
"""Process thing one by one, break the process chain when get
exception.
"""
self.logger.debug('Process thing: ' + str(thing))
raw_thing = thing
for pipeline in pipelines:
try:
thing = pipeline.process(thing)
if asyncio.iscoroutine(thing):
thing = await thing
except Exception as e:
if isinstance(e, ThingDropped):
self.report(raw_thing, dropped=True)
raise e
return thing
async def _request(self, req: Request) -> Response:
if req.proxy is not None:
# proxy auth not work in one session with many requests,
# add auth header to fix it
if req.proxy.scheme == 'http' and req.proxy.user is not None:
req.headers[aiohttp.hdrs.PROXY_AUTHORIZATION] = \
aiohttp.BasicAuth.from_url(req.proxy).encode()
# cookies in headers, params in url
req_kwargs = dict(headers=req.headers,
data=req.data,
timeout=req.timeout,
proxy=req.proxy,
max_redirects=self.request_max_redirects,
allow_redirects=self.request_allow_redirects)
response = await self.session._request(req.method, req.url,
**req_kwargs)
if not req.response_in_stream:
await response.read()
response.close()
await response.wait_for_close()
return response
class Client:
""" Client class handles all the connection and data transfer that happens
through the websocket for all types of image viewers.
"""
def __init__(self, addr, images, labels, **view_args):
"""
Args:
addr: address of the websocket
images: a list of images
labels: a list of labels for each image
"""
self._addr = addr
self._images = images
self._labels = labels
self.ws = None
# initialize queue and worker process
self.queue = Queue()
asyncio.create_task(self.worker())
asyncio.create_task(self.websocket_connect(view_args))
async def worker(self):
""" Handles asynchronous tasks that are enqueued to the work queue.
"""
while True:
message = await self.queue.get()
await self.ws.write_message(message)
async def websocket_connect(self, view_args):
""" Connects to the websocket at given address. Sends attributes and
data following a signal to the client.
Arguments:
view_args: a dictionary of arguments to be passed to the
Javascript side
"""
self.ws = await ws.websocket_connect(self._addr)
message = {"py_client": None}
await self.write_message(message)
# Wait to send data until signal received from server
message = await self.ws.read_message()
await self.write_message({"attrs": view_args})
await self.send_data()
async def write_message(self, message):
""" Writes message to the websocket.
Arguments:
message: a dictionary that contains the message to be sent to
the Javascript side
"""
message_json = json.dumps(message)
self.queue.put_nowait(message_json)
async def send_data(self):
""" Send image data to front end.
"""
images_bytes = list(map(Client.image2bytes, self._images))
data_dict = {"data": {}}
for i, image in enumerate(images_bytes):
data_dict["data"][i] = image
await self.write_message(data_dict)
@staticmethod
def image2bytes(image, quality=80):
""" Converts PIL Image to bytearray.
Arguments:
image: image buffer in PIL Image format
quality: the quality of the output image
"""
bytesIO = io.BytesIO()
image.save(bytesIO, format="JPEG", optimize=True, quality=quality)
return base64.b64encode(bytesIO.getvalue()).decode()
class HomeworksProtocol(asyncio.Protocol):
_non_login_reply_received_timer: TimerHandle
read_queue: Queue[Message]
_transport: Transport
PROMPT_REQUESTS = [b'LNET> ', b'L232> ']
LOGIN_REQUEST = b'LOGIN: '******'\r\n'
def __init__(self, credentials: Optional[Union[str, bytes]] = None):
self.ready_future = asyncio.Future()
self.connection_lost_future = asyncio.Future()
self.read_queue = Queue()
self._buffer = b''
self._credentials = ensure_bytes(credentials)
def data_received(self, data: bytes) -> None:
self._buffer += data
self.handle_buffer_increment()
def connection_made(self, transport: Transport) -> None:
self._transport = transport
self._non_login_reply_received_timer = asyncio.get_event_loop(
).call_later(0.2, self._notify_ready)
def connection_lost(self, exc: Optional[Exception]) -> None:
if not self._transport.is_closing():
self._transport.close()
self._transport = None
self._non_login_reply_received_timer.cancel()
exception = HomeworksConnectionLost(
f'Connection lost before ready state: {exc}')
if not self.ready_future.done():
self.ready_future.set_exception(exception)
self.connection_lost_future.set_exception(exception)
def handle_buffer_increment(self):
while any([
self._check_login_prompt(),
self._trim_prompts(),
self._check_messages()
]):
pass
def write(self, data: bytes):
data = ensure_bytes(data)
if not self._transport.is_closing():
self._transport.write(data)
def _check_login_prompt(self) -> bool:
return self._trim_prefix(self.LOGIN_REQUEST,
self._on_login_prompt_found)
def _trim_prompts(self) -> bool:
return any(
self._trim_prefix(prompt, self._on_prompt_found)
for prompt in self.PROMPT_REQUESTS)
def _on_prompt_found(self, _):
self._notify_ready()
def _on_login_prompt_found(self, _):
self._non_login_reply_received_timer.cancel()
if not self._credentials:
self._raise_exception(HomeworksNoCredentialsProvided())
self.write(self._credentials + self.COMMAND_SEPARATOR)
def _trim_prefix(self, prefix: bytes, on_match: Callable[[bytes],
None]) -> bool:
if self._buffer.startswith(prefix):
self._buffer = self._buffer[len(prefix):]
on_match(prefix)
return True
return False
def _check_messages(self) -> bool:
(command, separator,
remainder) = self._buffer.partition(self.COMMAND_SEPARATOR)
if separator != self.COMMAND_SEPARATOR:
return False
self._buffer = remainder
command = command.strip()
if command == b'':
return True
self._handle_message(command.decode(ENCODING))
return True
def _handle_message(self, message: str):
if message == "login successful":
self._notify_ready()
return
elif message == "login incorrect":
self._raise_exception(InvalidCredentialsProvided())
self._notify_ready()
self.read_queue.put_nowait(message)
def _notify_ready(self):
self._non_login_reply_received_timer.cancel()
if self._transport is not None and not self.ready_future.done():
self.ready_future.set_result(True)
def _raise_exception(self, exc: Exception):
if not self.ready_future.done():
self.ready_future.set_exception(exc)
raise exc
class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
"""
This class implements common parts of the WebSocket protocol.
It assumes that the WebSocket connection is established. The handshake is
managed in subclasses such as
:class:`~websockets.server.WebSocketServerProtocol` and
:class:`~websockets.client.WebSocketClientProtocol`.
It runs a task that stores incoming data frames in a queue and deals with
control frames automatically. It sends outgoing data frames and performs
the closing handshake.
The `host`, `port` and `secure` parameters are simply stored as attributes
for handlers that need them.
The `timeout` parameter defines the maximum wait time in seconds for
completing the closing handshake and, only on the client side, for
terminating the TCP connection. :meth:`close()` will complete in at most
this time on the server side and twice this time on the client side.
Once the connection is closed, the status code is available in the
:attr:`close_code` attribute and the reason in :attr:`close_reason`.
"""
# There are only two differences between the client-side and the server-
# side behavior: masking the payload and closing the underlying TCP
# connection. This class implements the server-side behavior by default.
# To get the client-side behavior, set is_client = True.
is_client = False
state = 'OPEN'
def __init__(self,
*,
host=None,
port=None,
secure=None,
timeout=10,
loop=None):
self.host = host
self.port = port
self.secure = secure
self.timeout = timeout
super().__init__(asyncio.StreamReader(), self.client_connected, loop)
self.close_code = None
self.close_reason = ''
# Futures tracking steps in the connection's lifecycle.
self.opening_handshake = asyncio.Future()
self.closing_handshake = asyncio.Future()
self.connection_closed = asyncio.Future()
# Queue of received messages.
self.messages = Queue()
# Mapping of ping IDs to waiters, in chronological order.
self.pings = collections.OrderedDict()
# Task managing the connection.
self.worker = asyncio. async (self.run())
# In a subclass implementing the opening handshake, the state will be
# CONNECTING at this point.
if self.state == 'OPEN':
self.opening_handshake.set_result(True)
# Public API
@property
def open(self):
"""
This property is ``True`` when the connection is usable.
It may be used to handle disconnections gracefully.
"""
return self.state == 'OPEN'
@asyncio.coroutine
def close(self, code=1000, reason=''):
"""
This coroutine performs the closing handshake.
This is the expected way to terminate a connection on the server side.
It waits for the other end to complete the handshake. It doesn't do
anything once the connection is closed.
It's usually safe to wrap this coroutine in `asyncio.async()` since
errors during connection termination aren't particularly useful.
The `code` must be an :class:`int` and the `reason` a :class:`str`.
"""
if self.state == 'OPEN':
# 7.1.2. Start the WebSocket Closing Handshake
self.close_code, self.close_reason = code, reason
yield from self.write_frame(OP_CLOSE,
serialize_close(code, reason))
# 7.1.3. The WebSocket Closing Handshake is Started
self.state = 'CLOSING'
# If the connection doesn't terminate within the timeout, break out of
# the worker loop.
try:
yield from asyncio.wait_for(self.worker, timeout=self.timeout)
except asyncio.TimeoutError:
self.worker.cancel()
# The worker should terminate quickly once it has been cancelled.
yield from self.worker
@asyncio.coroutine
def recv(self):
"""
This coroutine receives the next message.
It returns a :class:`str` for a text frame and :class:`bytes` for a
binary frame.
When the end of the message stream is reached, or when a protocol
error occurs, :meth:`recv` returns ``None``, indicating that the
connection is closed.
"""
# Return any available message
try:
return self.messages.get_nowait()
except QueueEmpty:
pass
# Wait for a message until the connection is closed
next_message = asyncio.Task(self.messages.get())
done, pending = yield from asyncio.wait(
[next_message, self.worker], return_when=asyncio.FIRST_COMPLETED)
if next_message in done:
return next_message.result()
@asyncio.coroutine
def send(self, data):
"""
This coroutine sends a message.
It sends a :class:`str` as a text frame and :class:`bytes` as a binary
frame.
It raises a :exc:`TypeError` for other inputs and
:exc:`InvalidState` once the connection is closed.
"""
if isinstance(data, str):
opcode = 1
data = data.encode('utf-8')
elif isinstance(data, bytes):
opcode = 2
else:
raise TypeError("data must be bytes or str")
yield from self.write_frame(opcode, data)
@asyncio.coroutine
def ping(self, data=None):
"""
This coroutine sends a ping.
It returns a Future which will be completed when the corresponding
pong is received and which you may ignore if you don't want to wait.
A ping may serve as a keepalive.
"""
# Protect against duplicates if a payload is explicitly set.
if data in self.pings:
raise ValueError("Already waiting for a pong with the same data")
# Generate a unique random payload otherwise.
while data is None or data in self.pings:
data = struct.pack('!I', random.getrandbits(32))
self.pings[data] = asyncio.Future()
yield from self.write_frame(OP_PING, data)
return self.pings[data]
@asyncio.coroutine
def pong(self, data=b''):
"""
This coroutine sends a pong.
An unsolicited pong may serve as a unidirectional heartbeat.
"""
yield from self.write_frame(OP_PONG, data)
# Private methods - no guarantees.
@asyncio.coroutine
def run(self):
# This coroutine guarantees that the connection is closed at exit.
yield from self.opening_handshake
while not self.closing_handshake.done():
try:
msg = yield from self.read_message()
if msg is None:
break
self.messages.put_nowait(msg)
except asyncio.CancelledError:
break
except WebSocketProtocolError:
yield from self.fail_connection(1002)
except UnicodeDecodeError:
yield from self.fail_connection(1007)
except Exception:
yield from self.fail_connection(1011)
raise
yield from self.close_connection()
@asyncio.coroutine
def read_message(self):
# Reassemble fragmented messages.
frame = yield from self.read_data_frame()
if frame is None:
return
if frame.opcode == OP_TEXT:
text = True
elif frame.opcode == OP_BINARY:
text = False
else: # frame.opcode == OP_CONT
raise WebSocketProtocolError("Unexpected opcode")
# Shortcut for the common case - no fragmentation
if frame.fin:
return frame.data.decode('utf-8') if text else frame.data
# 5.4. Fragmentation
chunks = []
if text:
decoder = codecs.getincrementaldecoder('utf-8')(errors='strict')
append = lambda f: chunks.append(decoder.decode(f.data, f.fin))
else:
append = lambda f: chunks.append(f.data)
append(frame)
while not frame.fin:
frame = yield from self.read_data_frame()
if frame is None:
raise WebSocketProtocolError("Incomplete fragmented message")
if frame.opcode != OP_CONT:
raise WebSocketProtocolError("Unexpected opcode")
append(frame)
return ('' if text else b'').join(chunks)
@asyncio.coroutine
def read_data_frame(self):
# Deal with control frames automatically and return next data frame.
# 6.2. Receiving Data
while True:
frame = yield from self.read_frame()
# 5.5. Control Frames
if frame.opcode == OP_CLOSE:
self.close_code, self.close_reason = parse_close(frame.data)
if self.state != 'CLOSING':
# 7.1.3. The WebSocket Closing Handshake is Started
self.state = 'CLOSING'
yield from self.write_frame(OP_CLOSE, frame.data,
'CLOSING')
self.closing_handshake.set_result(True)
return
elif frame.opcode == OP_PING:
# Answer pings.
yield from self.pong(frame.data)
elif frame.opcode == OP_PONG:
# Do not acknowledge pings on unsolicited pongs.
if frame.data in self.pings:
# Acknowledge all pings up to the one matching this pong.
ping_id = None
while ping_id != frame.data:
ping_id, waiter = self.pings.popitem(0)
if not waiter.cancelled():
waiter.set_result(None)
# 5.6. Data Frames
else:
return frame
@asyncio.coroutine
def read_frame(self):
is_masked = not self.is_client
frame = yield from read_frame(self.reader.readexactly, is_masked)
side = 'client' if self.is_client else 'server'
logger.debug("%s << %s", side, frame)
return frame
@asyncio.coroutine
def write_frame(self, opcode, data=b'', expected_state='OPEN'):
# This may happen if a user attempts to write on a closed connection.
if self.state != expected_state:
raise InvalidState("Cannot write to a WebSocket "
"in the {} state".format(self.state))
frame = Frame(True, opcode, data)
side = 'client' if self.is_client else 'server'
logger.debug("%s >> %s", side, frame)
is_masked = self.is_client
write_frame(frame, self.writer.write, is_masked)
# Handle flow control automatically.
try:
yield from self.writer.drain()
except ConnectionResetError:
pass
@asyncio.coroutine
def close_connection(self):
# 7.1.1. Close the WebSocket Connection
if self.state == 'CLOSED':
return
# Defensive assertion for protocol compliance.
if self.state != 'CLOSING': # pragma: no cover
raise InvalidState("Cannot close a WebSocket connection "
"in the {} state".format(self.state))
if self.is_client:
try:
yield from asyncio.wait_for(self.connection_closed,
timeout=self.timeout)
except (asyncio.CancelledError, asyncio.TimeoutError):
pass
if self.state == 'CLOSED':
return
if self.writer.can_write_eof():
self.writer.write_eof()
self.writer.close()
try:
yield from asyncio.wait_for(self.connection_closed,
timeout=self.timeout)
except (asyncio.CancelledError, asyncio.TimeoutError):
pass
@asyncio.coroutine
def fail_connection(self, code=1011, reason=''):
# Losing the connection usually results in a protocol error.
# Preserve the original error code in this case.
if self.close_code != 1006:
self.close_code, self.close_reason = code, reason
# 7.1.7. Fail the WebSocket Connection
logger.info("Failing the WebSocket connection: %d %s", code, reason)
if self.state == 'OPEN':
yield from self.write_frame(OP_CLOSE,
serialize_close(code, reason))
self.state = 'CLOSING'
if not self.closing_handshake.done():
self.closing_handshake.set_result(False)
yield from self.close_connection()
# asyncio StreamReaderProtocol methods
def client_connected(self, reader, writer):
self.reader = reader
self.writer = writer
def connection_lost(self, exc):
# 7.1.4. The WebSocket Connection is Closed
self.state = 'CLOSED'
if not self.connection_closed.done():
self.connection_closed.set_result(None)
if self.close_code is None:
self.close_code = 1006
super().connection_lost(exc)
class DCCClient:
def __init__(self, loop, network, send=False):
self.loop = loop
self.network = network
self.read_queue = Queue(loop=loop)
self.send = send #ugh what if i want to RECEIVE though.
#not sure what the use case would be but...?
@asyncio.coroutine
def connect(self, port=None):
if not self.send:
server = self.current_server = self.network.servers[0]
self._reader, self._writer = yield from server.connect(self.loop)
self._read_loop_task = asyncio.Task(self._start_read_loop())
asyncio.async(self._read_loop_task, loop=self.loop)
else:
self._waiting = asyncio.Lock()
yield from self._waiting.acquire()
if port:
self.network = yield from asyncio.start_server(self._handle_client,
host=socket.gethostbyname(socket.gethostname()), port=port, loop=self.loop)
else:
logger.error("No port provided for send")
@asyncio.coroutine
def _handle_client(self, client_reader, client_writer):
self._reader = client_reader
self._writer = client_writer
self._waiting.release()
self._read_loop_task = asyncio.Task(self._start_read_loop())
asyncio.async(self._read_loop_task, loop=self.loop)
@asyncio.coroutine
def disconnect(self):
yield from self._writer.drain()
self._writer.write_eof()
self._read_loop_task.cancel()
yield from self._read_loop_task
while not self._reader.at_eof():
yield from self._reader.readline()
if self.send:
self.network.close()
@asyncio.coroutine
def _start_read_loop(self):
if not self.send: #acks don't really do anything so don't listen for them
while not self._reader.at_eof():
try:
yield from self._read_message()
except CancelledError:
return
except Exception:
logger.exception("Smothering exception in DCC read loop")
@asyncio.coroutine
def _read_message(self):
line = yield from self._reader.readline()
m = re.match(b'(.*)(\r|\n|\r\n)$', line)
assert m
line = m.group(1)
message = DCCMessage.parse(line)
logger.debug("recv: %r", message)
event = DirectMessage(self, message)
self.read_queue.put_nowait((message, event))
@asyncio.coroutine
def read_event(self):
message, event = yield from self.read_queue.get()
return event
@asyncio.coroutine
def say(self, message, target=None, no_respond=None):
self.send_message(message)
@asyncio.coroutine
def send_message(self, message):
message = DCCMessage(message)
logger.debug("sent: %r", message)
self._writer.write(message.render().encode('utf8') + b'\r\n')
@asyncio.coroutine
def transfer(self, path):
yield from self._waiting.acquire()
f = open(str(path), 'rb')
block = b'\x01'
while block != b'':
block = f.read(1024)
self._writer.write(block)
f.close()
self._waiting.release()
return True
class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
"""
This class implements common parts of the WebSocket protocol.
It assumes that the WebSocket connection is established. The handshake is
managed in subclasses such as
:class:`~websockets.server.WebSocketServerProtocol` and
:class:`~websockets.client.WebSocketClientProtocol`.
It runs a task that stores incoming data frames in a queue and deals with
control frames automatically. It sends outgoing data frames and performs
the closing handshake.
The `host`, `port` and `secure` parameters are simply stored as attributes
for handlers that need them.
The `timeout` parameter defines the maximum wait time in seconds for
completing the closing handshake and, only on the client side, for
terminating the TCP connection. :meth:`close()` will complete in at most
this time on the server side and twice this time on the client side.
The `max_size` parameter enforces the maximum size for incoming messages
in bytes. The default value is 1MB. ``None`` disables the limit. If a
message larger than the maximum size is received, :meth:`recv()` will
return ``None`` and the connection will be closed with status code 1009.
Once the connection is closed, the status code is available in the
:attr:`close_code` attribute and the reason in :attr:`close_reason`.
"""
# There are only two differences between the client-side and the server-
# side behavior: masking the payload and closing the underlying TCP
# connection. This class implements the server-side behavior by default.
# To get the client-side behavior, set is_client = True.
is_client = False
state = 'OPEN'
def __init__(self, *,
host=None, port=None, secure=None, timeout=10, max_size=2 ** 20, loop=None):
self.host = host
self.port = port
self.secure = secure
self.timeout = timeout
self.max_size = max_size
super().__init__(asyncio.StreamReader(), self.client_connected, loop)
self.close_code = None
self.close_reason = ''
# Futures tracking steps in the connection's lifecycle.
self.opening_handshake = asyncio.Future()
self.closing_handshake = asyncio.Future()
self.connection_failed = asyncio.Future()
self.connection_closed = asyncio.Future()
# Queue of received messages.
self.messages = Queue()
# Mapping of ping IDs to waiters, in chronological order.
self.pings = collections.OrderedDict()
# Task managing the connection.
self.worker = asyncio.async(self.run())
# In a subclass implementing the opening handshake, the state will be
# CONNECTING at this point.
if self.state == 'OPEN':
self.opening_handshake.set_result(True)
# Public API
@property
def open(self):
"""
This property is ``True`` when the connection is usable.
It may be used to handle disconnections gracefully.
"""
return self.state == 'OPEN'
@asyncio.coroutine
def close(self, code=1000, reason=''):
"""
This coroutine performs the closing handshake.
This is the expected way to terminate a connection on the server side.
It waits for the other end to complete the handshake. It doesn't do
anything once the connection is closed.
It's usually safe to wrap this coroutine in `asyncio.async()` since
errors during connection termination aren't particularly useful.
The `code` must be an :class:`int` and the `reason` a :class:`str`.
"""
if self.state == 'OPEN':
# 7.1.2. Start the WebSocket Closing Handshake
self.close_code, self.close_reason = code, reason
yield from self.write_frame(OP_CLOSE, serialize_close(code, reason))
# 7.1.3. The WebSocket Closing Handshake is Started
self.state = 'CLOSING'
# If the connection doesn't terminate within the timeout, break out of
# the worker loop.
try:
yield from asyncio.wait_for(self.worker, timeout=self.timeout)
except asyncio.TimeoutError:
self.worker.cancel()
# The worker should terminate quickly once it has been cancelled.
yield from self.worker
@asyncio.coroutine
def recv(self):
"""
This coroutine receives the next message.
It returns a :class:`str` for a text frame and :class:`bytes` for a
binary frame.
When the end of the message stream is reached, or when a protocol
error occurs, :meth:`recv` returns ``None``, indicating that the
connection is closed.
"""
# Return any available message
try:
return self.messages.get_nowait()
except QueueEmpty:
pass
# Wait for a message until the connection is closed
next_message = asyncio.async(self.messages.get())
done, pending = yield from asyncio.wait(
[next_message, self.worker],
return_when=asyncio.FIRST_COMPLETED)
if next_message in done:
return next_message.result()
else:
next_message.cancel()
@asyncio.coroutine
def send(self, data):
"""
This coroutine sends a message.
It sends a :class:`str` as a text frame and :class:`bytes` as a binary
frame.
It raises a :exc:`TypeError` for other inputs and
:exc:`InvalidState` once the connection is closed.
"""
if isinstance(data, str):
opcode = 1
data = data.encode('utf-8')
elif isinstance(data, bytes):
opcode = 2
else:
raise TypeError("data must be bytes or str")
yield from self.write_frame(opcode, data)
@asyncio.coroutine
def ping(self, data=None):
"""
This coroutine sends a ping.
It returns a Future which will be completed when the corresponding
pong is received and which you may ignore if you don't want to wait.
A ping may serve as a keepalive.
"""
# Protect against duplicates if a payload is explicitly set.
if data in self.pings:
raise ValueError("Already waiting for a pong with the same data")
# Generate a unique random payload otherwise.
while data is None or data in self.pings:
data = struct.pack('!I', random.getrandbits(32))
self.pings[data] = asyncio.Future()
yield from self.write_frame(OP_PING, data)
return self.pings[data]
@asyncio.coroutine
def pong(self, data=b''):
"""
This coroutine sends a pong.
An unsolicited pong may serve as a unidirectional heartbeat.
"""
yield from self.write_frame(OP_PONG, data)
# Private methods - no guarantees.
@asyncio.coroutine
def run(self):
# This coroutine guarantees that the connection is closed at exit.
yield from self.opening_handshake
while not self.closing_handshake.done():
try:
msg = yield from self.read_message()
if msg is None:
break
self.messages.put_nowait(msg)
except asyncio.CancelledError:
break
except WebSocketProtocolError:
yield from self.fail_connection(1002)
except asyncio.IncompleteReadError:
yield from self.fail_connection(1006)
except UnicodeDecodeError:
yield from self.fail_connection(1007)
except PayloadTooBig:
yield from self.fail_connection(1009)
except Exception:
yield from self.fail_connection(1011)
raise
yield from self.close_connection()
@asyncio.coroutine
def read_message(self):
# Reassemble fragmented messages.
frame = yield from self.read_data_frame(max_size=self.max_size)
if frame is None:
return
if frame.opcode == OP_TEXT:
text = True
elif frame.opcode == OP_BINARY:
text = False
else: # frame.opcode == OP_CONT
raise WebSocketProtocolError("Unexpected opcode")
# Shortcut for the common case - no fragmentation
if frame.fin:
return frame.data.decode('utf-8') if text else frame.data
# 5.4. Fragmentation
chunks = []
max_size = self.max_size
if text:
decoder = codecs.getincrementaldecoder('utf-8')(errors='strict')
if max_size is None:
def append(frame):
nonlocal chunks
chunks.append(decoder.decode(frame.data, frame.fin))
else:
def append(frame):
nonlocal chunks, max_size
chunks.append(decoder.decode(frame.data, frame.fin))
max_size -= len(frame.data)
else:
if max_size is None:
def append(frame):
nonlocal chunks
chunks.append(frame.data)
else:
def append(frame):
nonlocal chunks, max_size
chunks.append(frame.data)
max_size -= len(frame.data)
append(frame)
while not frame.fin:
frame = yield from self.read_data_frame(max_size=max_size)
if frame is None:
raise WebSocketProtocolError("Incomplete fragmented message")
if frame.opcode != OP_CONT:
raise WebSocketProtocolError("Unexpected opcode")
append(frame)
return ('' if text else b'').join(chunks)
@asyncio.coroutine
def read_data_frame(self, max_size):
# Deal with control frames automatically and return next data frame.
# 6.2. Receiving Data
while True:
frame = yield from self.read_frame(max_size)
# 5.5. Control Frames
if frame.opcode == OP_CLOSE:
self.close_code, self.close_reason = parse_close(frame.data)
if self.state != 'CLOSING':
# 7.1.3. The WebSocket Closing Handshake is Started
self.state = 'CLOSING'
yield from self.write_frame(OP_CLOSE, frame.data, 'CLOSING')
if not self.closing_handshake.done():
self.closing_handshake.set_result(True)
return
elif frame.opcode == OP_PING:
# Answer pings.
yield from self.pong(frame.data)
elif frame.opcode == OP_PONG:
# Do not acknowledge pings on unsolicited pongs.
if frame.data in self.pings:
# Acknowledge all pings up to the one matching this pong.
ping_id = None
while ping_id != frame.data:
ping_id, waiter = self.pings.popitem(0)
if not waiter.cancelled():
waiter.set_result(None)
# 5.6. Data Frames
else:
return frame
@asyncio.coroutine
def read_frame(self, max_size):
is_masked = not self.is_client
frame = yield from read_frame(self.reader.readexactly, is_masked, max_size=max_size)
side = 'client' if self.is_client else 'server'
logger.debug("%s << %s", side, frame)
return frame
@asyncio.coroutine
def write_frame(self, opcode, data=b'', expected_state='OPEN'):
# This may happen if a user attempts to write on a closed connection.
if self.state != expected_state:
raise InvalidState("Cannot write to a WebSocket "
"in the {} state".format(self.state))
frame = Frame(True, opcode, data)
side = 'client' if self.is_client else 'server'
logger.debug("%s >> %s", side, frame)
is_masked = self.is_client
write_frame(frame, self.writer.write, is_masked)
try:
# Handle flow control automatically.
yield from self.writer.drain()
except ConnectionResetError:
# Terminate the connection if the socket died,
# unless it's already being closed.
if expected_state != 'CLOSING':
self.state = 'CLOSING'
yield from self.fail_connection(1006)
@asyncio.coroutine
def close_connection(self):
# 7.1.1. Close the WebSocket Connection
if self.state == 'CLOSED':
return
# Defensive assertion for protocol compliance.
if self.state != 'CLOSING': # pragma: no cover
raise InvalidState("Cannot close a WebSocket connection "
"in the {} state".format(self.state))
if self.is_client:
try:
yield from asyncio.wait_for(self.connection_closed,
timeout=self.timeout)
except (asyncio.CancelledError, asyncio.TimeoutError):
pass
if self.state == 'CLOSED':
return
# Attempt to terminate the TCP connection properly.
# If the socket is already closed, this will crash.
try:
if self.writer.can_write_eof():
self.writer.write_eof()
except Exception:
pass
self.writer.close()
try:
yield from asyncio.wait_for(self.connection_closed,
timeout=self.timeout)
except (asyncio.CancelledError, asyncio.TimeoutError):
pass
@asyncio.coroutine
def fail_connection(self, code=1011, reason=''):
# Avoid calling fail_connection more than once to minimize
# the consequences of race conditions between the two sides.
if self.connection_failed.done():
# Wait until the other coroutine calls connection_lost.
yield from self.connection_closed
return
else:
self.connection_failed.set_result(None)
# Losing the connection usually results in a protocol error.
# Preserve the original error code in this case.
if self.close_code != 1006:
self.close_code, self.close_reason = code, reason
# 7.1.7. Fail the WebSocket Connection
logger.info("Failing the WebSocket connection: %d %s", code, reason)
if self.state == 'OPEN':
yield from self.write_frame(OP_CLOSE, serialize_close(code, reason))
self.state = 'CLOSING'
if not self.closing_handshake.done():
self.closing_handshake.set_result(False)
yield from self.close_connection()
# asyncio StreamReaderProtocol methods
def client_connected(self, reader, writer):
self.reader = reader
self.writer = writer
def connection_lost(self, exc):
# 7.1.4. The WebSocket Connection is Closed
self.state = 'CLOSED'
if not self.connection_closed.done():
self.connection_closed.set_result(None)
if self.close_code is None:
self.close_code = 1006
super().connection_lost(exc)
class IRCClient:
"""Higher-level IRC client. Takes care of most of the hard parts of IRC:
incoming server messages are bundled into more intelligible events (see
``dywypi.event``), and commands that expect replies are implemented as
coroutines.
"""
def __init__(self, loop, network):
self.loop = loop
self.network = network
self.joined_channels = {} # name => Channel
# IRC server features, as reported by ISUPPORT, with defaults taken
# from the RFC.
self.len_nick = 9
self.len_channel = 200
self.len_message = 510
# These lengths don't have limits mentioned in the RFC, so going with
# the smallest known values in the wild
self.len_kick = 80
self.len_topic = 80
self.len_away = 160
self.max_watches = 0
self.max_targets = 1
self.channel_types = set('#&')
self.channel_modes = {} # TODO, haha.
self.channel_prefixes = {} # TODO here too. IRCMode is awkward.
self.network_title = self.network.name
self.features = {}
# Various intermediate state used for waiting for replies and
# aggregating multi-part replies
# TODO hmmm so what happens if state just gets left here forever? do
# we care?
self._pending_names = {}
self._names_futures = {}
self._pending_topics = {}
self._join_futures = {}
self.event_queue = Queue(loop=loop)
def get_channel(self, channel_name):
"""Returns a `Channel` object containing everything the client
definitively knows about the given channel.
Note that if you, say, ask for the topic of a channel you aren't in and
then immediately call `get_channel`, the returned object won't have its
topic populated. State is only tracked persistently for channels the
bot is in; otherwise there's no way to know whether or not it's stale.
"""
if channel_name in self.joined_channels:
return self.joined_channels[channel_name]
else:
return IRCChannel(self, channel_name)
@asyncio.coroutine
def connect(self):
"""Coroutine for connecting to a single server.
Note that this will nonblock until the client is "registered", defined
as the first PING/PONG exchange.
"""
# TODO this is a poor excuse for round-robin :)
server = self.current_server = self.network.servers[0]
# TODO i'm pretty sure the server tells us what our nick is, and we
# should believe that instead
self.nick = self.network.preferred_nick
# TODO: handle disconnection, somehow. probably affects a lot of
# things.
# TODO kind of wish this weren't here, since the creation of the
# connection isn't inherently part of a client. really it should be on
# the... network, perhaps? and there's no reason i shouldn't be able
# to "connect" to a unix socket or pipe or anywhere else that has data.
_, self.proto = yield from self.loop.create_connection(
lambda: IRCClientProtocol(
self.loop, self.network.preferred_nick, password=server.password),
server.host, server.port, ssl=server.tls)
while True:
yield from self._read_message()
# TODO this is dumb garbage; more likely this client itself should
# just wait for 001/RPL_WELCOME.
if self.proto.registered:
break
# Start the event loop as soon as we've synched, or we can't respond to
# anything
asyncio.async(self._advance(), loop=self.loop)
# Initial joins
yield from asyncio.gather(*[
self.join(channel_name)
for channel_name in self.network.autojoins
], loop=self.loop)
@asyncio.coroutine
def disconnect(self):
self.proto.send_message('QUIT', 'Seeya!')
self.proto.transport.close()
@asyncio.coroutine
def _advance(self):
"""Internal coroutine that just keeps the protocol message queue going.
Called once after a connect and should never be called again after
that.
"""
# TODO this is currently just to keep the message queue going, but
# eventually it should turn them into events and stuff them in an event
# queue
yield from self._read_message()
asyncio.async(self._advance(), loop=self.loop)
@asyncio.coroutine
def _read_message(self):
"""Internal dispatcher for messages received from the protocol."""
message = yield from self.proto.read_message()
# TODO there is a general ongoing problem here with matching up
# responses. ESPECIALLY when error codes are possible. something here
# is gonna have to get a bit fancier. maybe it should live at the
# protocol level, actually...?
# Boy do I ever hate this pattern but it's slightly more maintainable
# than a 500-line if tree.
handler = getattr(self, '_handle_' + message.command, None)
if handler:
handler(message)
def _handle_RPL_ISUPPORT(self, message):
me, *features, human_text = message.args
for feature_string in features:
feature, _, value = feature_string.partition('=')
if value is None:
value = True
self.features[feature] = value
if feature == 'NICKLEN':
self.len_nick = int(value)
elif feature == 'CHANNELLEN':
self.len_channel = int(value)
elif feature == 'KICKLEN':
self.len_kick = int(value)
elif feature == 'TOPICLEN':
self.len_topic = int(value)
elif feature == 'AWAYLEN':
self.len_away = int(value)
elif feature == 'WATCH':
self.max_watches = int(value)
elif feature == 'CHANTYPES':
self.channel_types = set(value)
elif feature == 'PREFIX':
# List of channel user modes, in relative priority order, in
# the format (ov)@+
assert value[0] == '('
letters, symbols = value[1:].split(')')
assert len(letters) == len(symbols)
self.channel_prefixes.clear()
for letter, symbol in zip(letters, symbols):
mode = IRCMode(letter, prefix=symbol)
self.channel_modes[letter] = mode
self.channel_prefixes[symbol] = mode
elif feature == 'MAXTARGETS':
self.max_targets = int(value)
elif feature == 'CHANMODES':
# Four groups delimited by lists: list-style (+b), arg required
# (+k), arg required only to set (+l), argless
lists, args, argsets, argless = value.split(',')
for letter in lists:
self.channel_modes[letter] = IRCMode(
letter, multi=True)
for letter in args:
self.channel_modes[letter] = IRCMode(
letter, arg_on_set=True, arg_on_remove=True)
for letter in argsets:
self.channel_modes[letter] = IRCMode(
letter, arg_on_set=True)
for letter in argless:
self.channel_modes[letter] = IRCMode(letter)
elif feature == 'NETWORK':
self.network_title = value
def _handle_JOIN(self, message):
channel_name, = message.args
joiner = Peer.from_prefix(message.prefix)
# TODO should there be a self.me? how...
if joiner.name == self.nick:
# We just joined a channel
#assert channel_name not in self.joined_channels
# TODO key? do we care?
# TODO what about channel configuration and anon non-joined
# channels? how do these all relate...
channel = IRCChannel(self, channel_name)
self.joined_channels[channel.name] = channel
else:
# Someone else just joined the channel
self.joined_channels[channel_name].add_user(joiner)
def _handle_RPL_TOPIC(self, message):
# Topic. Sent when joining or when requesting the topic.
# TODO this doesn't handle the "requesting" part
# TODO what if me != me?
me, channel_name, topic_text = message.args
self._pending_topics[channel_name] = IRCTopic(topic_text)
def _handle_RPL_TOPICWHOTIME(self, message):
# Topic author (NONSTANDARD). Sent after RPL_TOPIC.
# Unfortunately, there's no way to know whether to expect this.
# TODO this doesn't handle the "requesting" part
# TODO what if me != me?
me, channel_name, author, timestamp = message.args
topic = self._pending_topics.setdefault(channel_name, IRCTopic(''))
topic.author = Peer.from_prefix(author)
topic.timestamp = datetime.utcfromtimestamp(int(timestamp))
def _handle_RPL_NAMREPLY(self, message):
# Names response. Sent when joining or when requesting a names
# list. Must be ended with a RPL_ENDOFNAMES.
me, useless_equals_sign, channel_name, *raw_names = message.args
# List of names is actually optional (?!)
if raw_names:
raw_names = raw_names[0]
else:
raw_names = ''
names = raw_names.strip(' ').split(' ')
namelist = self._pending_names.setdefault(channel_name, [])
# TODO modes? should those be stripped off here?
# TODO for that matter should these become peers here?
namelist.extend(names)
def _handle_RPL_ENDOFNAMES(self, message):
# End of names list. Sent at the very end of a join or the very
# end of a NAMES request.
me, channel_name, info = message.args
namelist = self._pending_names.pop(channel_name, [])
if channel_name in self._names_futures:
# TODO we should probably not ever have a names future AND a
# pending join at the same time. or, does it matter?
self._names_futures[channel_name].set_result(namelist)
del self._names_futures[channel_name]
if channel_name in self.joined_channels:
# Join synchronized!
channel = self.joined_channels[channel_name]
channel.sync = True
channel.topic = self._pending_topics.pop(channel_name, None)
for name in namelist:
modes = set()
# TODO use features!
while name and name[0] in '+%@&~':
modes.add(name[0])
name = name[1:]
# TODO haha no this is so bad.
# TODO the bot should, obviously, keep a record of all
# known users as well. alas, mutable everything.
peer = Peer(name, None, None)
channel.add_user(peer, modes)
if channel_name in self._join_futures:
# Update the Future
self._join_futures[channel_name].set_result(channel)
del self._join_futures[channel_name]
def _handle_PRIVMSG(self, message):
event = Message(self, message)
self.event_queue.put_nowait(event)
@asyncio.coroutine
def read_event(self):
"""Produce a single IRC event.
This client does not do any kind of multiplexing or event handler
notification; that's left to a higher level.
"""
return (yield from self.event_queue.get())
# Implementations of particular commands
# TODO should this be part of the general client interface, or should there
# be a separate thing that smooths out the details?
@asyncio.coroutine
def say(self, target, message):
"""Coroutine that sends a message to a target, which may be either a
`Channel` or a `Peer`.
"""
yield from self.send_message('PRIVMSG', target, message)
def join(self, channel_name, key=None):
"""Coroutine that joins a channel, and nonblocks until the join is
"synchronized" (defined as receiving the nick list).
"""
if channel_name in self._join_futures:
return self._join_futures[channel_name]
# TODO multiple? error on commas?
if key is None:
self.proto.send_message('JOIN', channel_name)
else:
self.proto.send_message('JOIN', channel_name, key)
# Clear out any lingering names list
self._pending_names[channel_name] = []
# Return a Future, to be populated by the message loop
fut = self._join_futures[channel_name] = asyncio.Future()
return fut
def names(self, channel_name):
"""Coroutine that returns a list of names in a channel."""
self.proto.send_message('NAMES', channel_name)
# No need to do the same thing twice
if channel_name in self._names_futures:
return self._names_futures[channel_name]
# Clear out any lingering names list
self._pending_names[channel_name] = []
# Return a Future, to be populated by the message loop
fut = self._names_futures[channel_name] = asyncio.Future()
return fut
def set_topic(self, channel, topic):
"""Sets the channel topic."""
self.proto.send_message('TOPIC', channel, topic)
@asyncio.coroutine
def send_message(self, command, *args):
self.proto.send_message(command, *args)
def format_transition(self, current_style, new_style):
if new_style == Style.default():
# Reset code, ^O
return '\x0f'
if new_style.fg != current_style.fg and new_style.fg is Color.default:
# IRC has no "reset to default" code. mIRC claims color 99 is for
# this, but it lies, at least in irssi. So we must reset and
# reapply everything.
ret = '\x0f'
if new_style.bold is Bold.on:
ret += '\x02'
return ret
ret = ''
if new_style.fg != current_style.fg:
ret += FOREGROUND_CODES[new_style.fg]
if new_style.bold != current_style.bold:
# There's no on/off for bold, just a toggle
ret += '\x02'
return ret
async def dump1090TCPListener(incomingQ: Queue, host: str, port: int):
LOG.info(F"Starting: Opening TCP Stream to Dump Server at: {host}:{port}")
while True:
try:
stream_reader, _ = await open_connection(host, port)
STATS['connection_closed'] = False
STATS['connections'] += 1
except Exception as x:
print(F"@dump1090TCPListener: Got exception {x}")
LOG.exception(x)
STATS['connections_exception'] += 1
continue
LOG.info("Got stream_reader")
try:
BUFSIZE = 1024
data = ""
while True:
new_data = await stream_reader.read(BUFSIZE)
if not new_data:
LOG.info("receiver: connection closed")
STATS['connection_closed'] = True
return
# break it into lines
data += new_data.decode('utf-8')
#print(F"#######\n{data}######")
cur_pos = 0
while True:
# print(F"DATA\n{data}\n")
cr_index = data.find('\n')
if cr_index > 0:
# print(F"Slice {cur_pos} : {cr_index} Data: {data[cur_pos:cr_index]}")
l = data[cur_pos:cr_index].strip()
#print(F"Q<-{l}\n", end='')
incomingQ.put_nowait(l)
STATS["msg_rx"] += 1
STATS["max_incomingQ_size"] = max(
incomingQ.qsize(), STATS["max_incomingQ_size"])
# print(F"done")
if STATS["msg_rx"] % 100 == 0:
LOG.info(
F"Rxed {STATS['msg_rx']} max_q size {STATS['max_incomingQ_size']}"
)
cur_pos = cr_index + 1
data = data[cur_pos:]
cur_pos = 0
else:
# print(F"Slice {cur_pos} : {cr_index} Data: {data[cur_pos:cr_index]}")
break
except CancelledError:
LOG.info("Cancelling")
break
except Exception:
LOG.exception("General exception")
finally:
pass
LOG.info(F"Exiting")
class HandlerConnectionHandler:
def __init__(self, host, port, _handler_id, _device_ids, log, app,
event) -> None:
self._mqtt = None
self._app = app
self._log = log
self.event = None
self.command = None
self._client_id = f'handler.{_handler_id}'
self._host = host
self._port = port
self._device_ids = _device_ids
self._handler_id = _handler_id
self._event = event
self._message_queue = Queue()
def start(self):
self._mqtt = ConnectionHandler(self._host, self._port, self._client_id,
self._log)
self._log.set_mqtt_client(self._mqtt)
self._mqtt.init_mqtt_client_callbacks(self._on_connect,
self._on_message,
self._on_disconnect)
self._mqtt.register_route(
"Master", lambda topic, payload: self.dispatch_masterapp_message(
topic, self._mqtt.decode_payload(payload)))
self._mqtt.register_route(
"Handler", lambda topic, payload: self.dispatch_masterapp_message(
topic, self._mqtt.decode_payload(payload)))
self._mqtt.set_last_will(
self._generate_handler_status_topic(),
self._mqtt.create_message(
self._generate_status_message('crash', '')))
self._mqtt.start_loop()
async def stop(self) -> None:
await self._mqtt.stop_loop()
def subscribe(self, topic) -> None:
self._mqtt.subscribe(topic)
def publish(self, topic, payload, qos=0, retain=False) -> None:
self._mqtt.publish(topic, json.dumps(payload), qos=qos, retain=retain)
def publish_state(self, state, message) -> None:
self._log.log_message(LogLevel.Info(), f'Handler state: {state}')
self.publish(self._generate_handler_status_topic(),
self._generate_status_message(state, message),
qos=1,
retain=False)
def _on_connect(self, client, userdata, flags, rc) -> None:
self._app.startup_done('connection to broker is established')
for device_id in self._device_ids:
self.subscribe(self._generate_master_status_topic(device_id))
self.subscribe(self._generate_command_topic(device_id))
self.subscribe(self._generate_response_topic(device_id))
def _on_message(self, client, userdata, message) -> None:
self._mqtt.router.inject_message(message.topic, message.payload)
def _on_disconnect(self, client, userdata, rc) -> None:
self._log.debug("disconnected")
def dispatch_masterapp_message(self, topic, message) -> None:
master_id = topic.split('/')[1]
if "status" in topic:
master_state = message['state']
self._app.on_master_state_changed(master_id, master_state)
self._handle_master_state(message)
elif "command" in topic or\
"response" in topic:
self._put_message(message)
else:
assert False
def _handle_master_state(self, message):
if message['state'] in INVISIBLE_TESTER_STATES:
return
if not self.command:
return
self._put_message(message)
def _put_message(self, message):
self._message_queue.put_nowait(message)
self._event.set()
def get_message(self):
try:
return self._message_queue.get_nowait()
except QueueEmpty:
return None
def get_mqtt_client(self):
return self._mqtt
def send_command_message(self, message):
for device_id in self._device_ids:
self._mqtt.publish(self._generate_master_command_topic(device_id),
self._mqtt.create_message(message), False)
def send_response_message(self, message):
self._mqtt.publish(self._generate_handler_response_topic(),
self._mqtt.create_message(message), False)
def handle_message(self, message):
response = self._app.handle_message(message)
if len(response):
self._message_queue.put_nowait(response)
self._event.set()
return
if message['type'] == 'temperature':
self.send_response_message(message)
else:
self.send_command_message(message)
self.command = message['type']
@staticmethod
def _generate_message(type, payload):
return {'type': type, 'payload': payload}
def _generate_status_message(self, state, message):
payload = {'state': state, 'message': message}
return self._generate_message('status', payload)
@staticmethod
def _generate_command_topic(device_id):
return f"ate/{device_id}/Handler/command"
@staticmethod
def _generate_response_topic(device_id):
return f"ate/{device_id}/Master/response"
@staticmethod
def _generate_log_message(log_message):
return {"type": "log", "payload": log_message}
def _handler_log_topic(self):
return f'ate/{self._handler_id}/Handler/log/'
def _generate_handler_status_topic(self):
return f'ate/{self._handler_id}/Handler/status'
def _generate_handler_response_topic(self):
return f'ate/{self._handler_id}/Handler/response'
@staticmethod
def _generate_master_status_topic(device_id):
return f'ate/{device_id}/Master/status'
@staticmethod
def _generate_master_command_topic(device_id):
return f"ate/{device_id}/Master/command"
class Connection:
def __init__(self, cursor, conn: Database):
self._cursor = cursor
self._public_cursor = None
self._connection = conn
self._thread = ThreadPoolExecutor(max_workers=1)
self._result_queue = Queue(maxsize=1)
def _callback(self, result):
res = result.result()
self._result_queue.put_nowait(res)
async def __get_result(self):
while True:
try:
result = self._result_queue.get_nowait()
except QueueEmpty:
await asyncio.sleep(0.0001)
except Exception as e:
raise e
else:
return result
def _execute(self, sql: str, parameters=None):
if not parameters:
parameters = []
try:
result = self._cursor.execute(sql, parameters)
self._public_cursor = result
return self
except Exception as e:
return e
async def execute(self, sql, parameters=None):
task = self._thread.submit(self._execute, sql, parameters=parameters)
task.add_done_callback(self._callback)
result = await self.__get_result()
if isinstance(result, Exception):
raise result
return result
def _execute_many(self, sql, parameters):
try:
result = self._cursor.executemany(sql, parameters)
self._public_cursor = result
return self
except Exception as e:
return e
async def executemany(self, sql, parameters):
task = self._thread.submit(self._execute_many,
sql,
parameters=parameters)
task.add_done_callback(self._callback)
result = await self.__get_result()
if isinstance(result, Exception):
raise result
return result
def _fetch_one(self):
try:
result = self._public_cursor.fetchone()
return result
except Exception as e:
return e
async def fetchone(self):
if self._public_cursor is None:
raise NoActiveCursor
task = self._thread.submit(self._fetch_one)
task.add_done_callback(self._callback)
result = await self.__get_result()
if isinstance(result, Exception):
raise result
return result
def _fetch_all(self):
try:
result = self._public_cursor.fetchall()
return result
except Exception as e:
return e
async def fetchall(self):
if self._public_cursor is None:
raise NoActiveCursor
task = self._thread.submit(self._fetch_all)
task.add_done_callback(self._callback)
result = await self.__get_result()
if isinstance(result, Exception):
raise result
return result
def _commit(self):
try:
self._cursor.commit()
return None
except Exception as e:
return e
async def commit(self):
task = self._thread.submit(self._commit)
task.add_done_callback(self._callback)
result = await self.__get_result()
if isinstance(result, Exception):
raise result
return result
async def close(self):
self._cursor.close()
await self._connection.connection_semaphore.put({})
return
