class Collector(BaseHandler):
"""
The Collector handler allows for collecting a set of stanzas
that match a given pattern. Unlike the Waiter handler, a
Collector does not block execution, and will continue to
accumulate matching stanzas until told to stop.
:param string name: The name of the handler.
:param matcher: A :class:`~sleekxmpp.xmlstream.matcher.base.MatcherBase`
derived object for matching stanza objects.
:param stream: The :class:`~sleekxmpp.xmlstream.xmlstream.XMLStream`
instance this handler should monitor.
"""
def __init__(self, name, matcher, stream=None):
BaseHandler.__init__(self, name, matcher, stream=stream)
self._payload = Queue()
def prerun(self, payload):
"""Store the matched stanza when received during processing.
:param payload: The matched
:class:`~sleekxmpp.xmlstream.stanzabase.ElementBase` object.
"""
self._payload.put(payload)
def run(self, payload):
"""Do not process this handler during the main event loop."""
pass
def stop(self):
"""
Stop collection of matching stanzas, and return the ones that
have been stored so far.
"""
self._destroy = True
results = []
try:
while True:
results.append(self._payload.get(False))
except QueueEmpty:
pass
self.stream().remove_handler(self.name)
return results
class bot(sleekxmpp.ClientXMPP):
'''
Bot used to send and receive messages
'''
def __init__(self, master, jid_password):
self.master = master
self.buffer_size = 0 #number of bytes of data waiting to be sent
self.buffer_queue = Queue()
self.buffer_size_lock = threading.Lock()
self.num_clients = 0 #number of clients sockets using this bot
self.num_clients_lock = threading.Lock()
self.__failed_send_stanza = None #some sleekxmpp thing for resending stuff
sleekxmpp.ClientXMPP.__init__(self, *jid_password)
# gmail xmpp server is actually at talk.google.com
if jid_password[0].find("@gmail.com") != -1:
self.connect_address = ("talk.google.com", 5222)
else:
self.connect_address = None
# event handlers are sleekxmpp's way of dealing with important xml tags it receives
self.add_event_handler("session_start",
lambda event: self.session_start())
self.add_event_handler("disconnected",
lambda event: self.disconnected())
self.register_plugin('xep_0199') # Ping
def buffer_message(self, data):
'''
set the karma and time_last_sent
karma_lock should have been acquired before calling this function
'''
self.buffer_queue.put(data)
self.buffer_size += len(data)
self.buffer_size_lock.release()
def get_buffer_size(self):
self.buffer_size_lock.acquire()
return self.buffer_size
def get_num_clients(self):
self.num_clients_lock.acquire()
return self.num_clients
def register_hexchat_handlers(self):
'''these handle the custom iq stanzas'''
self.register_handler(
callback.Callback('Connect Handler',
stanzapath.StanzaPath('iq@type=set/connect'),
self.master.connect_handler))
self.register_handler(
callback.Callback(
'Connect Message Handler',
stanzapath.StanzaPath('message@type=chat/connect'),
self.master.connect_handler))
self.register_handler(
callback.Callback(
'Connect Ack Handler',
stanzapath.StanzaPath('iq@type=result/connect_ack'),
self.master.connect_ack_handler))
self.register_handler(
callback.Callback('Data Handler',
stanzapath.StanzaPath('iq@type=set/packet'),
self.master.data_handler))
self.register_handler(
callback.Callback('Disconnect Handler',
stanzapath.StanzaPath('iq@type=set/disconnect'),
self.master.disconnect_handler))
self.register_handler(
callback.Callback(
'Disconnect Error Message Handler',
stanzapath.StanzaPath('message@type=chat/disconnect_error'),
self.master.disconnect_error_handler))
self.register_handler(
callback.Callback(
'Disconnect Error Iq Handler',
stanzapath.StanzaPath('iq@type=set/disconnect_error'),
self.master.disconnect_error_handler))
self.register_handler(
callback.Callback('IQ Error Handler',
stanzapath.StanzaPath('iq@type=error/error'),
self.master.error_handler))
self.register_handler(
callback.Callback(
'Message Error Handler',
stanzapath.StanzaPath('message@type=error/error'),
self.master.error_handler))
### session management mathods:
def boot(self, process=True):
if self.connect(self.connect_address):
if process:
self.process()
else:
raise (Exception(self.boundjid.bare + " could not connect"))
def session_start(self):
"""Called when the bot connects and establishes a session with the XMPP server."""
# XMPP spec says that we should broadcast our presence when we connect.
self.send_presence()
def disconnected(self):
'''
Called when the bot disconnects from the XMPP server.
Try to reconnect.
'''
logging.warning("XMPP chat server disconnected")
logging.debug("Trying to reconnect")
self.boot(False)
self.send_presence()
def _send_thread(self):
'''
modifed version of sleekxmpp's _send_thread
that will not send faster than THROUGHPUT
'''
try:
while not self.stop.is_set():
while not self.stop.is_set() and \
not self.session_started_event.is_set():
self.session_started_event.wait(timeout=0.1)
if self.__failed_send_stanza is not None:
data = self.__failed_send_stanza
self.__failed_send_stanza = None
else:
try:
data = self.buffer_queue.get(True, 1)
was_buffered = True
except QueueEmpty:
try:
data = self.send_queue.get(True, 0.0)
was_buffered = False
except QueueEmpty:
continue
logging.debug("SEND: %s", data)
enc_data = data.encode('utf-8')
total = len(enc_data)
sent = 0
count = 0
tries = 0
try:
with self.send_lock:
while sent < total and not self.stop.is_set() and \
self.session_started_event.is_set():
try:
num_bytes = self.socket.send(enc_data[sent:])
sent += num_bytes
count += 1
'''
throttling code
that prevents data from being sent
faster than THROUGHPUT
'''
time.sleep(num_bytes / THROUGHPUT)
if was_buffered:
with self.buffer_size_lock:
self.buffer_size -= num_bytes
except ssl.SSLError as serr:
if tries >= self.ssl_retry_max:
logging.debug(
'SSL error: max retries reached')
self.exception(serr)
logging.warning("Failed to send %s", data)
if not self.stop.is_set():
self.disconnect(self.auto_reconnect,
send_close=False)
logging.warning(
'SSL write error: retrying')
if not self.stop.is_set():
time.sleep(self.ssl_retry_delay)
tries += 1
if count > 1:
logging.debug('SENT: %d chunks', count)
if was_buffered:
self.buffer_queue.task_done()
else:
self.send_queue.task_done()
except (Socket.error, ssl.SSLError) as serr:
self.event('socket_error', serr, direct=True)
logging.warning("Failed to send %s", data)
if not self.stop.is_set():
self.__failed_send_stanza = data
self._end_thread('send')
self.disconnect(self.auto_reconnect, send_close=False)
return
except Exception as ex:
logging.exception('Unexpected error in send thread: %s', ex)
self.exception(ex)
if not self.stop.is_set():
self._end_thread('send')
self.disconnect(self.auto_reconnect)
return
self._end_thread('send')
class Waiter(BaseHandler):
"""
The Waiter handler allows an event handler to block until a
particular stanza has been received. The handler will either be
given the matched stanza, or ``False`` if the waiter has timed out.
:param string name: The name of the handler.
:param matcher: A :class:`~sleekxmpp.xmlstream.matcher.base.MatcherBase`
derived object for matching stanza objects.
:param stream: The :class:`~sleekxmpp.xmlstream.xmlstream.XMLStream`
instance this handler should monitor.
"""
def __init__(self, name, matcher, stream=None):
BaseHandler.__init__(self, name, matcher, stream=stream)
self._payload = Queue()
def prerun(self, payload):
"""Store the matched stanza when received during processing.
:param payload: The matched
:class:`~sleekxmpp.xmlstream.stanzabase.ElementBase` object.
"""
self._payload.put(payload)
def run(self, payload):
"""Do not process this handler during the main event loop."""
pass
def wait(self, timeout=None):
"""Block an event handler while waiting for a stanza to arrive.
Be aware that this will impact performance if called from a
non-threaded event handler.
Will return either the received stanza, or ``False`` if the
waiter timed out.
:param int timeout: The number of seconds to wait for the stanza
to arrive. Defaults to the the stream's
:class:`~sleekxmpp.xmlstream.xmlstream.XMLStream.response_timeout`
value.
"""
if timeout is None:
timeout = self.stream().response_timeout
elapsed_time = 0
stanza = False
while elapsed_time < timeout and not self.stream().stop.is_set():
try:
stanza = self._payload.get(True, 1)
break
except QueueEmpty:
elapsed_time += 1
if elapsed_time >= timeout:
log.warning("Timed out waiting for %s", self.name)
self.stream().remove_handler(self.name)
return stanza
def check_delete(self):
"""Always remove waiters after use."""
return True
class IBBytestream(object):
def __init__(self,
xmpp,
sid,
block_size,
to,
ifrom,
window_size=1,
use_messages=False):
self.xmpp = xmpp
self.sid = sid
self.block_size = block_size
self.window_size = window_size
self.use_messages = use_messages
self.receiver = to
self.sender = ifrom
self.send_seq = -1
self.recv_seq = -1
self._send_seq_lock = threading.Lock()
self._recv_seq_lock = threading.Lock()
self.stream_started = threading.Event()
self.stream_in_closed = threading.Event()
self.stream_out_closed = threading.Event()
self.recv_queue = Queue()
self.send_window = threading.BoundedSemaphore(value=self.window_size)
self.window_ids = set()
self.window_empty = threading.Event()
self.window_empty.set()
def send(self, data):
if not self.stream_started.is_set() or \
self.stream_out_closed.is_set():
raise socket.error
data = data[0:self.block_size]
self.send_window.acquire()
with self._send_seq_lock:
self.send_seq = (self.send_seq + 1) % 65535
seq = self.send_seq
if self.use_messages:
msg = self.xmpp.Message()
msg['to'] = self.receiver
msg['from'] = self.sender
msg['id'] = self.xmpp.new_id()
msg['ibb_data']['sid'] = self.sid
msg['ibb_data']['seq'] = seq
msg['ibb_data']['data'] = data
msg.send()
self.send_window.release()
else:
iq = self.xmpp.Iq()
iq['type'] = 'set'
iq['to'] = self.receiver
iq['from'] = self.sender
iq['ibb_data']['sid'] = self.sid
iq['ibb_data']['seq'] = seq
iq['ibb_data']['data'] = data
self.window_empty.clear()
self.window_ids.add(iq['id'])
iq.send(block=False, callback=self._recv_ack)
return len(data)
def sendall(self, data):
sent_len = 0
while sent_len < len(data):
sent_len += self.send(data[sent_len:])
def _recv_ack(self, iq):
self.window_ids.remove(iq['id'])
if not self.window_ids:
self.window_empty.set()
self.send_window.release()
if iq['type'] == 'error':
self.close()
def _recv_data(self, stanza):
with self._recv_seq_lock:
new_seq = stanza['ibb_data']['seq']
if new_seq != (self.recv_seq + 1) % 65535:
self.close()
raise XMPPError('unexpected-request')
self.recv_seq = new_seq
data = stanza['ibb_data']['data']
if len(data) > self.block_size:
self.close()
raise XMPPError('not-acceptable')
self.recv_queue.put(data)
self.xmpp.event('ibb_stream_data', {'stream': self, 'data': data})
if isinstance(stanza, Iq):
stanza.reply()
stanza.send()
def recv(self, *args, **kwargs):
return self.read(block=True)
def read(self, block=True, timeout=None, **kwargs):
if not self.stream_started.is_set() or \
self.stream_in_closed.is_set():
raise socket.error
if timeout is not None:
block = True
try:
return self.recv_queue.get(block, timeout)
except:
return None
def close(self):
iq = self.xmpp.Iq()
iq['type'] = 'set'
iq['to'] = self.receiver
iq['from'] = self.sender
iq['ibb_close']['sid'] = self.sid
self.stream_out_closed.set()
iq.send(block=False, callback=lambda x: self.stream_in_closed.set())
self.xmpp.event('ibb_stream_end', self)
def _closed(self, iq):
self.stream_in_closed.set()
self.stream_out_closed.set()
while not self.window_empty.is_set():
log.info('waiting for send window to empty')
self.window_empty.wait(timeout=1)
iq.reply()
iq.send()
self.xmpp.event('ibb_stream_end', self)
def makefile(self, *args, **kwargs):
return self
def connect(*args, **kwargs):
return None
def shutdown(self, *args, **kwargs):
return None
class IBBytestream(object):
def __init__(self, xmpp, sid, block_size, to, ifrom, window_size=1):
self.xmpp = xmpp
self.sid = sid
self.block_size = block_size
self.window_size = window_size
self.receiver = to
self.sender = ifrom
self.send_seq = -1
self.recv_seq = -1
self._send_seq_lock = threading.Lock()
self._recv_seq_lock = threading.Lock()
self.stream_started = threading.Event()
self.stream_in_closed = threading.Event()
self.stream_out_closed = threading.Event()
self.recv_queue = Queue()
self.send_window = threading.BoundedSemaphore(value=self.window_size)
self.window_ids = set()
self.window_empty = threading.Event()
self.window_empty.set()
def send(self, data):
if not self.stream_started.is_set() or \
self.stream_out_closed.is_set():
raise socket.error
data = data[0:self.block_size]
self.send_window.acquire()
with self._send_seq_lock:
self.send_seq = (self.send_seq + 1) % 65535
seq = self.send_seq
iq = self.xmpp.Iq()
iq['type'] = 'set'
iq['to'] = self.receiver
iq['from'] = self.sender
iq['ibb_data']['sid'] = self.sid
iq['ibb_data']['seq'] = seq
iq['ibb_data']['data'] = data
self.window_empty.clear()
self.window_ids.add(iq['id'])
iq.send(block=False, callback=self._recv_ack)
return len(data)
def sendall(self, data):
sent_len = 0
while sent_len < len(data):
sent_len += self.send(data[sent_len:])
def _recv_ack(self, iq):
self.window_ids.remove(iq['id'])
if not self.window_ids:
self.window_empty.set()
self.send_window.release()
if iq['type'] == 'error':
self.close()
def _recv_data(self, iq):
with self._recv_seq_lock:
new_seq = iq['ibb_data']['seq']
if new_seq != (self.recv_seq + 1) % 65535:
self.close()
raise XMPPError('unexpected-request')
self.recv_seq = new_seq
data = iq['ibb_data']['data']
if len(data) > self.block_size:
self.close()
raise XMPPError('not-acceptable')
self.recv_queue.put(data)
self.xmpp.event('ibb_stream_data', {'stream': self, 'data': data})
iq.reply()
iq.send()
def recv(self, *args, **kwargs):
return self.read(block=True)
def read(self, block=True, timeout=None, **kwargs):
if not self.stream_started.is_set() or \
self.stream_in_closed.is_set():
raise socket.error
if timeout is not None:
block = True
try:
return self.recv_queue.get(block, timeout)
except:
return None
def close(self):
iq = self.xmpp.Iq()
iq['type'] = 'set'
iq['to'] = self.receiver
iq['from'] = self.sender
iq['ibb_close']['sid'] = self.sid
self.stream_out_closed.set()
iq.send(block=False,
callback=lambda x: self.stream_in_closed.set())
self.xmpp.event('ibb_stream_end', self)
def _closed(self, iq):
self.stream_in_closed.set()
self.stream_out_closed.set()
while not self.window_empty.is_set():
log.info('waiting for send window to empty')
self.window_empty.wait(timeout=1)
iq.reply()
iq.send()
self.xmpp.event('ibb_stream_end', self)
def makefile(self, *args, **kwargs):
return self
def connect(*args, **kwargs):
return None
def shutdown(self, *args, **kwargs):
return None
class bot(sleekxmpp.ClientXMPP):
'''
Bot used to send and receive messages
'''
def __init__(self, master, jid_password):
self.master = master
self.buffer_size = 0 #number of bytes of data waiting to be sent
self.buffer_queue = Queue()
self.buffer_size_lock = threading.Lock()
self.num_clients = 0 #number of clients sockets using this bot
self.num_clients_lock = threading.Lock()
self.__failed_send_stanza = None #some sleekxmpp thing for resending stuff
sleekxmpp.ClientXMPP.__init__(self, *jid_password)
# gmail xmpp server is actually at talk.google.com
if jid_password[0].find("@gmail.com") != -1:
self.connect_address = ("talk.google.com", 5222)
else:
self.connect_address = None
# event handlers are sleekxmpp's way of dealing with important xml tags it receives
self.add_event_handler("session_start", lambda event: self.session_start())
self.add_event_handler("disconnected", lambda event: self.disconnected())
self.register_plugin('xep_0199') # Ping
def buffer_message(self, data):
'''
set the karma and time_last_sent
karma_lock should have been acquired before calling this function
'''
self.buffer_queue.put(data)
self.buffer_size += len(data)
self.buffer_size_lock.release()
def get_buffer_size(self):
self.buffer_size_lock.acquire()
return self.buffer_size
def get_num_clients(self):
self.num_clients_lock.acquire()
return self.num_clients
def register_hexchat_handlers(self):
'''these handle the custom iq stanzas'''
self.register_handler(
callback.Callback('Connect Handler',
stanzapath.StanzaPath('iq@type=set/connect'),
self.master.connect_handler)
)
self.register_handler(
callback.Callback('Connect Message Handler',
stanzapath.StanzaPath('message@type=chat/connect'),
self.master.connect_handler)
)
self.register_handler(
callback.Callback('Connect Ack Handler',
stanzapath.StanzaPath('iq@type=result/connect_ack'),
self.master.connect_ack_handler)
)
self.register_handler(
callback.Callback('Data Handler',
stanzapath.StanzaPath('iq@type=set/packet'),
self.master.data_handler)
)
self.register_handler(
callback.Callback('Disconnect Handler',
stanzapath.StanzaPath('iq@type=set/disconnect'),
self.master.disconnect_handler)
)
self.register_handler(
callback.Callback('Disconnect Error Message Handler',
stanzapath.StanzaPath('message@type=chat/disconnect_error'),
self.master.disconnect_error_handler)
)
self.register_handler(
callback.Callback('Disconnect Error Iq Handler',
stanzapath.StanzaPath('iq@type=set/disconnect_error'),
self.master.disconnect_error_handler)
)
self.register_handler(
callback.Callback('IQ Error Handler',
stanzapath.StanzaPath('iq@type=error/error'),
self.master.error_handler)
)
self.register_handler(
callback.Callback('Message Error Handler',
stanzapath.StanzaPath('message@type=error/error'),
self.master.error_handler)
)
### session management mathods:
def boot(self, process=True):
if self.connect(self.connect_address):
if process:
self.process()
else:
raise(Exception(self.boundjid.bare + " could not connect"))
def session_start(self):
"""Called when the bot connects and establishes a session with the XMPP server."""
# XMPP spec says that we should broadcast our presence when we connect.
self.send_presence()
def disconnected(self):
'''
Called when the bot disconnects from the XMPP server.
Try to reconnect.
'''
logging.warning("XMPP chat server disconnected")
logging.debug("Trying to reconnect")
self.boot(False)
self.send_presence()
def _send_thread(self):
'''
modifed version of sleekxmpp's _send_thread
that will not send faster than THROUGHPUT
'''
try:
while not self.stop.is_set():
while not self.stop.is_set() and \
not self.session_started_event.is_set():
self.session_started_event.wait(timeout=0.1)
if self.__failed_send_stanza is not None:
data = self.__failed_send_stanza
self.__failed_send_stanza = None
else:
try:
data = self.buffer_queue.get(True, 1)
was_buffered = True
except QueueEmpty:
try:
data = self.send_queue.get(True, 0.0)
was_buffered = False
except QueueEmpty:
continue
logging.debug("SEND: %s", data)
enc_data = data.encode('utf-8')
total = len(enc_data)
sent = 0
count = 0
tries = 0
try:
with self.send_lock:
while sent < total and not self.stop.is_set() and \
self.session_started_event.is_set():
try:
num_bytes = self.socket.send(enc_data[sent:])
sent += num_bytes
count += 1
'''
throttling code
that prevents data from being sent
faster than THROUGHPUT
'''
time.sleep(num_bytes / THROUGHPUT)
if was_buffered:
with self.buffer_size_lock:
self.buffer_size -= num_bytes
except ssl.SSLError as serr:
if tries >= self.ssl_retry_max:
logging.debug('SSL error: max retries reached')
self.exception(serr)
logging.warning("Failed to send %s", data)
if not self.stop.is_set():
self.disconnect(self.auto_reconnect,
send_close=False)
logging.warning('SSL write error: retrying')
if not self.stop.is_set():
time.sleep(self.ssl_retry_delay)
tries += 1
if count > 1:
logging.debug('SENT: %d chunks', count)
if was_buffered:
self.buffer_queue.task_done()
else:
self.send_queue.task_done()
except (Socket.error, ssl.SSLError) as serr:
self.event('socket_error', serr, direct=True)
logging.warning("Failed to send %s", data)
if not self.stop.is_set():
self.__failed_send_stanza = data
self._end_thread('send')
self.disconnect(self.auto_reconnect, send_close=False)
return
except Exception as ex:
logging.exception('Unexpected error in send thread: %s', ex)
self.exception(ex)
if not self.stop.is_set():
self._end_thread('send')
self.disconnect(self.auto_reconnect)
return
self._end_thread('send')
class Scheduler(object):
"""
A threaded scheduler that allows for updates mid-execution unlike the
scheduler in the standard library.
Based on: http://docs.python.org/library/sched.html#module-sched
:param parentstop: An :class:`~threading.Event` to signal stopping
the scheduler.
"""
def __init__(self, parentstop=None):
#: A queue for storing tasks
self.addq = Queue()
#: A list of tasks in order of execution time.
self.schedule = []
#: If running in threaded mode, this will be the thread processing
#: the schedule.
self.thread = None
#: A flag indicating that the scheduler is running.
self.run = False
#: An :class:`~threading.Event` instance for signalling to stop
#: the scheduler.
self.stop = parentstop
#: Lock for accessing the task queue.
self.schedule_lock = threading.RLock()
def process(self, threaded=True, daemon=False):
"""Begin accepting and processing scheduled tasks.
:param bool threaded: Indicates if the scheduler should execute
in its own thread. Defaults to ``True``.
"""
if threaded:
self.thread = threading.Thread(name='scheduler_process',
target=self._process)
self.thread.daemon = daemon
self.thread.start()
else:
self._process()
def _process(self):
"""Process scheduled tasks."""
self.run = True
try:
while self.run and not self.stop.is_set():
wait = 0.1
updated = False
if self.schedule:
wait = self.schedule[0].next - time.time()
try:
if wait <= 0.0:
newtask = self.addq.get(False)
else:
if wait >= 3.0:
wait = 3.0
newtask = None
elapsed = 0
while not self.stop.is_set() and \
newtask is None and \
elapsed < wait:
newtask = self.addq.get(True, 0.1)
elapsed += 0.1
except QueueEmpty:
self.schedule_lock.acquire()
# select only those tasks which are to be executed now
relevant = itertools.takewhile(
lambda task: time.time() >= task.next, self.schedule)
# run the tasks and keep the return value in a tuple
status = map(lambda task: (task, task.run()), relevant)
# remove non-repeating tasks
for task, doRepeat in status:
if not doRepeat:
try:
self.schedule.remove(task)
except ValueError:
pass
else:
# only need to resort tasks if a repeated task has
# been kept in the list.
updated = True
else:
updated = True
self.schedule_lock.acquire()
if newtask is not None:
self.schedule.append(newtask)
finally:
if updated:
self.schedule.sort(key=lambda task: task.next)
self.schedule_lock.release()
except KeyboardInterrupt:
self.run = False
except SystemExit:
self.run = False
log.debug("Quitting Scheduler thread")
def add(self,
name,
seconds,
callback,
args=None,
kwargs=None,
repeat=False,
qpointer=None):
"""Schedule a new task.
:param string name: The name of the task.
:param int seconds: The number of seconds to wait before executing.
:param callback: The function to execute.
:param tuple args: The arguments to pass to the callback.
:param dict kwargs: The keyword arguments to pass to the callback.
:param bool repeat: Indicates if the task should repeat.
Defaults to ``False``.
:param pointer: A pointer to an event queue for queuing callback
execution instead of executing immediately.
"""
try:
self.schedule_lock.acquire()
for task in self.schedule:
if task.name == name:
raise ValueError("Key %s already exists" % name)
self.addq.put(
Task(name, seconds, callback, args, kwargs, repeat, qpointer))
except:
raise
finally:
self.schedule_lock.release()
def remove(self, name):
"""Remove a scheduled task ahead of schedule, and without
executing it.
:param string name: The name of the task to remove.
"""
try:
self.schedule_lock.acquire()
the_task = None
for task in self.schedule:
if task.name == name:
the_task = task
if the_task is not None:
self.schedule.remove(the_task)
except:
raise
finally:
self.schedule_lock.release()
def quit(self):
"""Shutdown the scheduler."""
self.run = False
class Scheduler(object):
"""
A threaded scheduler that allows for updates mid-execution unlike the
scheduler in the standard library.
Based on: http://docs.python.org/library/sched.html#module-sched
:param parentstop: An :class:`~threading.Event` to signal stopping
the scheduler.
"""
def __init__(self, parentstop=None):
#: A queue for storing tasks
self.addq = Queue()
#: A list of tasks in order of execution time.
self.schedule = []
#: If running in threaded mode, this will be the thread processing
#: the schedule.
self.thread = None
#: A flag indicating that the scheduler is running.
self.run = False
#: An :class:`~threading.Event` instance for signalling to stop
#: the scheduler.
self.stop = parentstop
#: Lock for accessing the task queue.
self.schedule_lock = threading.RLock()
def process(self, threaded=True, daemon=False):
"""Begin accepting and processing scheduled tasks.
:param bool threaded: Indicates if the scheduler should execute
in its own thread. Defaults to ``True``.
"""
if threaded:
self.thread = threading.Thread(name='scheduler_process',
target=self._process)
self.thread.daemon = daemon
self.thread.start()
else:
self._process()
def _process(self):
"""Process scheduled tasks."""
self.run = True
try:
while self.run and not self.stop.is_set():
wait = 0.1
updated = False
if self.schedule:
wait = self.schedule[0].next - time.time()
try:
if wait <= 0.0:
newtask = self.addq.get(False)
else:
if wait >= 3.0:
wait = 3.0
newtask = None
elapsed = 0
while not self.stop.is_set() and \
newtask is None and \
elapsed < wait:
newtask = self.addq.get(True, 0.1)
elapsed += 0.1
except QueueEmpty:
cleanup = []
self.schedule_lock.acquire()
for task in self.schedule:
if time.time() >= task.next:
updated = True
if not task.run():
cleanup.append(task)
else:
break
for task in cleanup:
self.schedule.pop(self.schedule.index(task))
else:
updated = True
self.schedule_lock.acquire()
if newtask is not None:
self.schedule.append(newtask)
finally:
if updated:
self.schedule = sorted(self.schedule,
key=lambda task: task.next)
self.schedule_lock.release()
except KeyboardInterrupt:
self.run = False
except SystemExit:
self.run = False
log.debug("Quitting Scheduler thread")
def add(self, name, seconds, callback, args=None,
kwargs=None, repeat=False, qpointer=None):
"""Schedule a new task.
:param string name: The name of the task.
:param int seconds: The number of seconds to wait before executing.
:param callback: The function to execute.
:param tuple args: The arguments to pass to the callback.
:param dict kwargs: The keyword arguments to pass to the callback.
:param bool repeat: Indicates if the task should repeat.
Defaults to ``False``.
:param pointer: A pointer to an event queue for queuing callback
execution instead of executing immediately.
"""
try:
self.schedule_lock.acquire()
for task in self.schedule:
if task.name == name:
raise ValueError("Key %s already exists" % name)
self.addq.put(Task(name, seconds, callback, args,
kwargs, repeat, qpointer))
except:
raise
finally:
self.schedule_lock.release()
def remove(self, name):
"""Remove a scheduled task ahead of schedule, and without
executing it.
:param string name: The name of the task to remove.
"""
try:
self.schedule_lock.acquire()
the_task = None
for task in self.schedule:
if task.name == name:
the_task = task
if the_task is not None:
self.schedule.remove(the_task)
except:
raise
finally:
self.schedule_lock.release()
def quit(self):
"""Shutdown the scheduler."""
self.run = False
class IBBytestream(object):
def __init__(self, xmpp, sid, block_size, jid, peer, window_size=1, use_messages=False):
self.xmpp = xmpp
self.sid = sid
self.block_size = block_size
self.window_size = window_size
self.use_messages = use_messages
if jid is None:
jid = xmpp.boundjid
self.self_jid = jid
self.peer_jid = peer
self.send_seq = -1
self.recv_seq = -1
self._send_seq_lock = threading.Lock()
self._recv_seq_lock = threading.Lock()
self.stream_started = threading.Event()
self.stream_in_closed = threading.Event()
self.stream_out_closed = threading.Event()
self.recv_queue = Queue()
self.send_window = threading.BoundedSemaphore(value=self.window_size)
self.window_ids = set()
self.window_empty = threading.Event()
self.window_empty.set()
def send(self, data):
if not self.stream_started.is_set() or \
self.stream_out_closed.is_set():
raise socket.error
data = data[0:self.block_size]
self.send_window.acquire()
with self._send_seq_lock:
self.send_seq = (self.send_seq + 1) % 65535
seq = self.send_seq
if self.use_messages:
msg = self.xmpp.Message()
msg['to'] = self.peer_jid
msg['from'] = self.self_jid
msg['id'] = self.xmpp.new_id()
msg['ibb_data']['sid'] = self.sid
msg['ibb_data']['seq'] = seq
msg['ibb_data']['data'] = data
msg.send()
self.send_window.release()
else:
iq = self.xmpp.Iq()
iq['type'] = 'set'
iq['to'] = self.peer_jid
iq['from'] = self.self_jid
iq['ibb_data']['sid'] = self.sid
iq['ibb_data']['seq'] = seq
iq['ibb_data']['data'] = data
self.window_empty.clear()
self.window_ids.add(iq['id'])
iq.send(block=False, callback=self._recv_ack)
return len(data)
def sendall(self, data):
sent_len = 0
while sent_len < len(data):
sent_len += self.send(data[sent_len:])
def _recv_ack(self, iq):
self.window_ids.remove(iq['id'])
if not self.window_ids:
self.window_empty.set()
self.send_window.release()
if iq['type'] == 'error':
self.close()
def _recv_data(self, stanza):
with self._recv_seq_lock:
new_seq = stanza['ibb_data']['seq']
if new_seq != (self.recv_seq + 1) % 65535:
self.close()
raise XMPPError('unexpected-request')
self.recv_seq = new_seq
data = stanza['ibb_data']['data']
if len(data) > self.block_size:
self.close()
raise XMPPError('not-acceptable')
self.recv_queue.put(data)
self.xmpp.event('ibb_stream_data', {'stream': self, 'data': data})
if isinstance(stanza, Iq):
stanza.reply()
stanza.send()
def recv(self, *args, **kwargs):
return self.read(block=True)
def read(self, block=True, timeout=None, **kwargs):
if not self.stream_started.is_set() or \
self.stream_in_closed.is_set():
raise socket.error
if timeout is not None:
block = True
try:
return self.recv_queue.get(block, timeout)
except:
return None
def close(self):
iq = self.xmpp.Iq()
iq['type'] = 'set'
iq['to'] = self.peer_jid
iq['from'] = self.self_jid
iq['ibb_close']['sid'] = self.sid
self.stream_out_closed.set()
iq.send(block=False,
callback=lambda x: self.stream_in_closed.set())
self.xmpp.event('ibb_stream_end', self)
def _closed(self, iq):
self.stream_in_closed.set()
self.stream_out_closed.set()
iq.reply()
iq.send()
self.xmpp.event('ibb_stream_end', self)
def makefile(self, *args, **kwargs):
return self
def connect(*args, **kwargs):
return None
def shutdown(self, *args, **kwargs):
return None
class TestLiveSocket(object):
"""
A live test socket that reads and writes to queues in
addition to an actual networking socket.
Methods:
next_sent -- Return the next sent stanza.
next_recv -- Return the next received stanza.
recv_data -- Dummy method to have same interface as TestSocket.
recv -- Read the next stanza from the socket.
send -- Write a stanza to the socket.
makefile -- Dummy call, returns self.
read -- Read the next stanza from the socket.
"""
def __init__(self, *args, **kwargs):
"""
Create a new, live test socket.
Arguments:
Same as arguments for socket.socket
"""
self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self.recv_buffer = []
self.recv_queue = Queue()
self.send_queue = Queue()
self.send_queue_lock = threading.Lock()
self.recv_queue_lock = threading.Lock()
self.is_live = True
def __getattr__(self, name):
"""
Return attribute values of internal, live socket.
Arguments:
name -- Name of the attribute requested.
"""
return getattr(self.socket, name)
# ------------------------------------------------------------------
# Testing Interface
def disconnect_errror(self):
"""
Used to simulate a socket disconnection error.
Not used by live sockets.
"""
try:
self.socket.shutdown()
self.socket.close()
except:
pass
def next_sent(self, timeout=None):
"""
Get the next stanza that has been sent.
Arguments:
timeout -- Optional timeout for waiting for a new value.
"""
args = {'block': False}
if timeout is not None:
args = {'block': True, 'timeout': timeout}
try:
return self.send_queue.get(**args)
except:
return None
def next_recv(self, timeout=None):
"""
Get the next stanza that has been received.
Arguments:
timeout -- Optional timeout for waiting for a new value.
"""
args = {'block': False}
if timeout is not None:
args = {'block': True, 'timeout': timeout}
try:
if self.recv_buffer:
return self.recv_buffer.pop(0)
else:
return self.recv_queue.get(**args)
except:
return None
def recv_data(self, data):
"""
Add data to a receive buffer for cases when more than a single stanza
was received.
"""
self.recv_buffer.append(data)
# ------------------------------------------------------------------
# Socket Interface
def recv(self, *args, **kwargs):
"""
Read data from the socket.
Store a copy in the receive queue.
Arguments:
Placeholders. Same as for socket.recv.
"""
data = self.socket.recv(*args, **kwargs)
with self.recv_queue_lock:
self.recv_queue.put(data)
return data
def send(self, data):
"""
Send data on the socket.
Store a copy in the send queue.
Arguments:
data -- String value to write.
"""
with self.send_queue_lock:
self.send_queue.put(data)
return self.socket.send(data)
# ------------------------------------------------------------------
# File Socket
def makefile(self, *args, **kwargs):
"""
File socket version to use with ElementTree.
Arguments:
Placeholders, same as socket.makefile()
"""
return self
def read(self, *args, **kwargs):
"""
Implement the file socket read interface.
Arguments:
Placeholders, same as socket.recv()
"""
return self.recv(*args, **kwargs)
def clear(self):
"""
Empty the send queue, typically done once the session has started to
remove the feature negotiation and log in stanzas.
"""
with self.send_queue_lock:
for i in range(0, self.send_queue.qsize()):
self.send_queue.get(block=False)
with self.recv_queue_lock:
for i in range(0, self.recv_queue.qsize()):
self.recv_queue.get(block=False)
class Scheduler(object):
"""
A threaded scheduler that allows for updates mid-execution unlike the
scheduler in the standard library.
Based on: http://docs.python.org/library/sched.html#module-sched
:param parentstop: An :class:`~threading.Event` to signal stopping
the scheduler.
"""
def __init__(self, parentstop=None):
#: A queue for storing tasks
self.addq = Queue()
#: A list of tasks in order of execution time.
self.schedule = []
#: If running in threaded mode, this will be the thread processing
#: the schedule.
self.thread = None
#: A flag indicating that the scheduler is running.
self.run = False
#: An :class:`~threading.Event` instance for signalling to stop
#: the scheduler.
self.stop = parentstop
#: Lock for accessing the task queue.
self.schedule_lock = threading.RLock()
#: The time in seconds to wait for events from the event queue,
#: and also the time between checks for the process stop signal.
self.wait_timeout = WAIT_TIMEOUT
def process(self, threaded=True, daemon=False):
"""Begin accepting and processing scheduled tasks.
:param bool threaded: Indicates if the scheduler should execute
in its own thread. Defaults to ``True``.
"""
if threaded:
self.thread = threading.Thread(name='scheduler_process',
target=self._process)
self.thread.daemon = daemon
self.thread.start()
else:
self._process()
def _process(self):
"""Process scheduled tasks."""
self.run = True
try:
while self.run and not self.stop.is_set():
updated = False
if self.schedule:
wait = self.schedule[0].next - time.time()
else:
wait = self.wait_timeout
try:
if wait <= 0.0:
newtask = self.addq.get(False)
else:
newtask = None
while self.run and \
not self.stop.is_set() and \
newtask is None and \
wait > 0:
try:
newtask = self.addq.get(True, min(wait, self.wait_timeout))
except QueueEmpty: # Nothing to add, nothing to do. Check run flags and continue waiting.
wait -= self.wait_timeout
except QueueEmpty: # Time to run some tasks, and no new tasks to add.
self.schedule_lock.acquire()
# select only those tasks which are to be executed now
relevant = itertools.takewhile(
lambda task: time.time() >= task.next, self.schedule)
# run the tasks and keep the return value in a tuple
status = map(lambda task: (task, task.run()), relevant)
# remove non-repeating tasks
for task, doRepeat in status:
if not doRepeat:
try:
self.schedule.remove(task)
except ValueError:
pass
else:
# only need to resort tasks if a repeated task has
# been kept in the list.
updated = True
else: # Add new task
self.schedule_lock.acquire()
if newtask is not None:
self.schedule.append(newtask)
updated = True
finally:
if updated:
self.schedule.sort(key=lambda task: task.next)
self.schedule_lock.release()
except KeyboardInterrupt:
self.run = False
except SystemExit:
self.run = False
log.debug("Quitting Scheduler thread")
def add(self, name, seconds, callback, args=None,
kwargs=None, repeat=False, qpointer=None):
"""Schedule a new task.
:param string name: The name of the task.
:param int seconds: The number of seconds to wait before executing.
:param callback: The function to execute.
:param tuple args: The arguments to pass to the callback.
:param dict kwargs: The keyword arguments to pass to the callback.
:param bool repeat: Indicates if the task should repeat.
Defaults to ``False``.
:param pointer: A pointer to an event queue for queuing callback
execution instead of executing immediately.
"""
try:
self.schedule_lock.acquire()
for task in self.schedule:
if task.name == name:
raise ValueError("Key %s already exists" % name)
self.addq.put(Task(name, seconds, callback, args,
kwargs, repeat, qpointer))
except:
raise
finally:
self.schedule_lock.release()
def remove(self, name):
"""Remove a scheduled task ahead of schedule, and without
executing it.
:param string name: The name of the task to remove.
"""
try:
self.schedule_lock.acquire()
the_task = None
for task in self.schedule:
if task.name == name:
the_task = task
if the_task is not None:
self.schedule.remove(the_task)
except:
raise
finally:
self.schedule_lock.release()
def quit(self):
"""Shutdown the scheduler."""
self.run = False
class TestSocket(object):
"""
A dummy socket that reads and writes to queues instead
of an actual networking socket.
Methods:
next_sent -- Return the next sent stanza.
recv_data -- Make a stanza available to read next.
recv -- Read the next stanza from the socket.
send -- Write a stanza to the socket.
makefile -- Dummy call, returns self.
read -- Read the next stanza from the socket.
"""
def __init__(self, *args, **kwargs):
"""
Create a new test socket.
Arguments:
Same as arguments for socket.socket
"""
self.socket = socket.socket(*args, **kwargs)
self.recv_queue = Queue()
self.send_queue = Queue()
self.is_live = False
self.disconnected = False
def __getattr__(self, name):
"""
Return attribute values of internal, dummy socket.
Some attributes and methods are disabled to prevent the
socket from connecting to the network.
Arguments:
name -- Name of the attribute requested.
"""
def dummy(*args):
"""Method to do nothing and prevent actual socket connections."""
return None
overrides = {'connect': dummy, 'close': dummy, 'shutdown': dummy}
return overrides.get(name, getattr(self.socket, name))
# ------------------------------------------------------------------
# Testing Interface
def next_sent(self, timeout=None):
"""
Get the next stanza that has been 'sent'.
Arguments:
timeout -- Optional timeout for waiting for a new value.
"""
args = {'block': False}
if timeout is not None:
args = {'block': True, 'timeout': timeout}
try:
return self.send_queue.get(**args)
except:
return None
def recv_data(self, data):
"""
Add data to the receiving queue.
Arguments:
data -- String data to 'write' to the socket to be received
by the XMPP client.
"""
self.recv_queue.put(data)
def disconnect_error(self):
"""
Simulate a disconnect error by raising a socket.error exception
for any current or further socket operations.
"""
self.disconnected = True
# ------------------------------------------------------------------
# Socket Interface
def recv(self, *args, **kwargs):
"""
Read a value from the received queue.
Arguments:
Placeholders. Same as for socket.Socket.recv.
"""
if self.disconnected:
raise socket.error
return self.read(block=True)
def send(self, data):
"""
Send data by placing it in the send queue.
Arguments:
data -- String value to write.
"""
if self.disconnected:
raise socket.error
self.send_queue.put(data)
return len(data)
# ------------------------------------------------------------------
# File Socket
def makefile(self, *args, **kwargs):
"""
File socket version to use with ElementTree.
Arguments:
Placeholders, same as socket.Socket.makefile()
"""
return self
def read(self, block=True, timeout=None, **kwargs):
"""
Implement the file socket interface.
Arguments:
block -- Indicate if the read should block until a
value is ready.
timeout -- Time in seconds a block should last before
returning None.
"""
if self.disconnected:
raise socket.error
if timeout is not None:
block = True
try:
return self.recv_queue.get(block, timeout)
except:
return None
class Superfeedr(ClientXMPP):
def __init__(self, jid, password, auto_connect=True):
ClientXMPP.__init__(self, jid, password)
self.success = False
self.notification_callback = None
self.wait_for_start = Queue()
self.register_plugin('xep_0004')
self.register_plugin('xep_0030')
self.register_plugin('xep_0060')
self.register_plugin('xep_0199')
self.add_event_handler("session_start", self.start)
handler = Callback('superfeedr',
MatchXPath("{jabber:client}message/"
"{http://jabber.org/protocol/pubsub#event}event"),
self.superfeedr_msg)
self.register_handler(handler)
if auto_connect:
self.connect_wait()
def connect_wait(self, timeout=10):
self.success = self.connect(('xmpp.superfeedr.com', 5222))
if self.success:
self.process(threaded=True)
start = self.wait_for_start.get(timeout)
if start is None:
self.success = False
def start(self, event):
self.get_roster()
self.send_presence()
self.wait_for_start.put(True)
def superfeedr_msg(self, stanza):
xml = stanza.xml
event = {'items': [],
'status': self.parse_status(xml, '{http://jabber.org/protocol/pubsub#event}event')}
items = xml.findall('{http://jabber.org/protocol/pubsub#event}event/'
'{http://jabber.org/protocol/pubsub#event}items/'
'{http://jabber.org/protocol/pubsub}item/'
'{http://www.w3.org/2005/Atom}content')
for item in items:
if item.get('type') == 'application/json':
event['items'].append(json.loads(item.text))
self.event('superfeedr', event)
if len(event.get('items', [])) > 0:
self.event('superfeedr_entry', event)
return event
def subscribe(self, feeds):
if len(feeds) > 20:
raise ValueError('Maximum of 20 feeds allowed per subscription message.')
pubsub = ElementTree.Element('{http://jabber.org/protocol/pubsub}pubsub')
pubsub.attrib['xmlns:superfeedr'] = 'http://superfeedr.com/xmpp-pubsub-ext'
for f in feeds:
feed = ElementTree.Element('subscribe')
feed.attrib['node'] = f
feed.attrib['jid'] = self.boundjid.bare
feed.attrib['superfeedr:format'] = 'json'
pubsub.append(feed)
iq = self.make_iq_set(pubsub)
iq.attrib['to'] = 'firehoser.superfeedr.com'
iq.attrib['from'] = self.boundjid.bare
iq.attrib['type'] = 'set'
response = self.send_wait(iq)
subscriptions = response.findall('{http://jabber.org/protocol/pubsub}pubsub/'
'{http://jabber.org/protocol/pubsub}subscription')
result = False
if subscriptions:
result = []
for subscription in subscriptions:
status = self.parse_status(subscription)
result.append({"subscription": {"status": status,
"feed": {"url": subscription.get('node'),
"title": status['title']}}})
return result
def unsubscribe(self, feed):
return self.plugin['xep_0060'].unsubscribe('firehoser.superfeedr.com', feed)
def list(self, page=0):
pubsub = ElementTree.Element('{http://jabber.org/protocol/pubsub}pubsub')
pubsub.attrib['xmlns:superfeedr'] = 'http://superfeedr.com/xmpp-pubsub-ext'
subscriptions = ElementTree.Element('subscriptions')
subscriptions.attrib['jid'] = self.boundjid.full
subscriptions.attrib['superfeedr:page'] = str(page)
pubsub.append(subscriptions)
iq = self.make_iq_set(pubsub)
iq.attrib['to'] = 'firehoser.superfeedr.com'
iq.attrib['from'] = self.boundjid.full
iq.attrib['type'] = 'get'
result = self.send_wait(iq)
if not result or result.get('type') == 'error':
return False
nodes = result.findall('{http://jabber.org/protocol/pubsub}pubsub/'
'{http://jabber.org/protocol/pubsub}subscriptions/'
'{http://jabber.org/protocol/pubsub}subscription')
if nodes is None:
return []
nodelist = []
for node in nodes:
nodelist.append(node.get('node', ''))
return nodelist
def on_notification(self, callback):
self.add_event_handler('superfeedr', callback)
def on_entry(self, callback):
self.add_event_handler('superfeedr_entry', callback)
def send_wait(self, iq, timeout=None):
"""
:param iq: Stanza to send
:param int timeout: The number of seconds to wait for the stanza
to arrive. Defaults to the the stream's
:class:`~sleekxmpp.xmlstream.xmlstream.XMLStream.response_timeout`
value.
"""
iq_id = iq.get('id')
self.send(iq)
waiter = Waiter("SendWait_%s" % self.new_id(), MatcherId(iq_id))
self.register_handler(waiter)
return waiter.wait(timeout)
@staticmethod
def parse_status(xml, base_query=None):
base_query = base_query + '/' if base_query else ''
status_query = base_query + '{http://superfeedr.com/xmpp-pubsub-ext}status'
status_child_query = status_query + '/{http://superfeedr.com/xmpp-pubsub-ext}'
status = xml.find(status_query)
http = xml.find(status_child_query + 'http')
next_fetch = xml.find(status_child_query + 'next_fetch')
last_fetch = xml.find(status_child_query + 'last_fetch')
last_parse = xml.find(status_child_query + 'last_parse')
period = xml.find(status_child_query + 'period')
last_maintenance_at = xml.find(status_child_query + 'last_maintenance_at')
title = xml.find(status_child_query + 'title')
if None not in (status, http, next_fetch, last_fetch, last_parse, period,
last_maintenance_at, last_fetch, title):
result = dict(lastParse=date_to_epoch(last_parse.text),
lastFetch=date_to_epoch(last_fetch.text),
nextFetch=date_to_epoch(next_fetch.text),
lastMaintenanceAt=date_to_epoch(last_maintenance_at.text),
period=period.text,
title=title.text)
feed = status.get('feed')
if feed is not None:
result['feed'] = feed
if http.text is not None:
result['http'] = http.text
code = http.get('code')
if code:
result['code'] = http.get('code')
return result
class Waiter(BaseHandler):
"""
The Waiter handler allows an event handler to block until a
particular stanza has been received. The handler will either be
given the matched stanza, or ``False`` if the waiter has timed out.
:param string name: The name of the handler.
:param matcher: A :class:`~sleekxmpp.xmlstream.matcher.base.MatcherBase`
derived object for matching stanza objects.
:param stream: The :class:`~sleekxmpp.xmlstream.xmlstream.XMLStream`
instance this handler should monitor.
"""
def __init__(self, name, matcher, stream=None):
BaseHandler.__init__(self, name, matcher, stream=stream)
self._payload = Queue()
def prerun(self, payload):
"""Store the matched stanza when received during processing.
:param payload: The matched
:class:`~sleekxmpp.xmlstream.stanzabase.ElementBase` object.
"""
self._payload.put(payload)
def run(self, payload):
"""Do not process this handler during the main event loop."""
pass
def wait(self, timeout=None):
"""Block an event handler while waiting for a stanza to arrive.
Be aware that this will impact performance if called from a
non-threaded event handler.
Will return either the received stanza, or ``False`` if the
waiter timed out.
:param int timeout: The number of seconds to wait for the stanza
to arrive. Defaults to the the stream's
:class:`~sleekxmpp.xmlstream.xmlstream.XMLStream.response_timeout`
value.
"""
if timeout is None:
timeout = self.stream().response_timeout
elapsed_time = 0
stanza = False
while elapsed_time < timeout and not self.stream().stop.is_set():
try:
stanza = self._payload.get(True, 1)
break
except QueueEmpty:
elapsed_time += 1
if elapsed_time >= timeout:
log.warning("Timed out waiting for %s", self.name)
self.stream().remove_handler(self.name)
return stanza
def check_delete(self):
"""Always remove waiters after use."""
return True
class TestSocket(object):
"""
A dummy socket that reads and writes to queues instead
of an actual networking socket.
Methods:
next_sent -- Return the next sent stanza.
recv_data -- Make a stanza available to read next.
recv -- Read the next stanza from the socket.
send -- Write a stanza to the socket.
makefile -- Dummy call, returns self.
read -- Read the next stanza from the socket.
"""
def __init__(self, *args, **kwargs):
"""
Create a new test socket.
Arguments:
Same as arguments for socket.socket
"""
self.socket = socket.socket(*args, **kwargs)
self.recv_queue = Queue()
self.send_queue = Queue()
self.is_live = False
self.disconnected = False
def __getattr__(self, name):
"""
Return attribute values of internal, dummy socket.
Some attributes and methods are disabled to prevent the
socket from connecting to the network.
Arguments:
name -- Name of the attribute requested.
"""
def dummy(*args):
"""Method to do nothing and prevent actual socket connections."""
return None
overrides = {'connect': dummy,
'close': dummy,
'shutdown': dummy}
return overrides.get(name, getattr(self.socket, name))
# ------------------------------------------------------------------
# Testing Interface
def next_sent(self, timeout=None):
"""
Get the next stanza that has been 'sent'.
Arguments:
timeout -- Optional timeout for waiting for a new value.
"""
args = {'block': False}
if timeout is not None:
args = {'block': True, 'timeout': timeout}
try:
return self.send_queue.get(**args)
except:
return None
def recv_data(self, data):
"""
Add data to the receiving queue.
Arguments:
data -- String data to 'write' to the socket to be received
by the XMPP client.
"""
self.recv_queue.put(data)
def disconnect_error(self):
"""
Simulate a disconnect error by raising a socket.error exception
for any current or further socket operations.
"""
self.disconnected = True
# ------------------------------------------------------------------
# Socket Interface
def recv(self, *args, **kwargs):
"""
Read a value from the received queue.
Arguments:
Placeholders. Same as for socket.Socket.recv.
"""
if self.disconnected:
raise socket.error
return self.read(block=True)
def send(self, data):
"""
Send data by placing it in the send queue.
Arguments:
data -- String value to write.
"""
if self.disconnected:
raise socket.error
self.send_queue.put(data)
return len(data)
# ------------------------------------------------------------------
# File Socket
def makefile(self, *args, **kwargs):
"""
File socket version to use with ElementTree.
Arguments:
Placeholders, same as socket.Socket.makefile()
"""
return self
def read(self, block=True, timeout=None, **kwargs):
"""
Implement the file socket interface.
Arguments:
block -- Indicate if the read should block until a
value is ready.
timeout -- Time in seconds a block should last before
returning None.
"""
if self.disconnected:
raise socket.error
if timeout is not None:
block = True
try:
return self.recv_queue.get(block, timeout)
except:
return None
