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 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
def stream_start(self,
mode='client',
skip=True,
header=None,
socket='mock',
jid='tester@localhost',
password='******',
server='localhost',
port=5222,
sasl_mech=None,
plugins=None,
plugin_config={}):
"""
Initialize an XMPP client or component using a dummy XML stream.
Arguments:
mode -- Either 'client' or 'component'. Defaults to 'client'.
skip -- Indicates if the first item in the sent queue (the
stream header) should be removed. Tests that wish
to test initializing the stream should set this to
False. Otherwise, the default of True should be used.
socket -- Either 'mock' or 'live' to indicate if the socket
should be a dummy, mock socket or a live, functioning
socket. Defaults to 'mock'.
jid -- The JID to use for the connection.
Defaults to 'tester@localhost'.
password -- The password to use for the connection.
Defaults to 'test'.
server -- The name of the XMPP server. Defaults to 'localhost'.
port -- The port to use when connecting to the server.
Defaults to 5222.
plugins -- List of plugins to register. By default, all plugins
are loaded.
"""
if mode == 'client':
self.xmpp = ClientXMPP(jid,
password,
sasl_mech=sasl_mech,
plugin_config=plugin_config)
elif mode == 'component':
self.xmpp = ComponentXMPP(jid,
password,
server,
port,
plugin_config=plugin_config)
else:
raise ValueError("Unknown XMPP connection mode.")
# Remove unique ID prefix to make it easier to test
self.xmpp._id_prefix = ''
self.xmpp._disconnect_wait_for_threads = False
self.xmpp.default_lang = None
self.xmpp.peer_default_lang = None
# We will use this to wait for the session_start event
# for live connections.
skip_queue = Queue()
if socket == 'mock':
self.xmpp.set_socket(TestSocket())
# Simulate connecting for mock sockets.
self.xmpp.auto_reconnect = False
self.xmpp.state._set_state('connected')
# Must have the stream header ready for xmpp.process() to work.
if not header:
header = self.xmpp.stream_header
self.xmpp.socket.recv_data(header)
elif socket == 'live':
self.xmpp.socket_class = TestLiveSocket
def wait_for_session(x):
self.xmpp.socket.clear()
skip_queue.put('started')
self.xmpp.add_event_handler('session_start', wait_for_session)
if server is not None:
self.xmpp.connect((server, port))
else:
self.xmpp.connect()
else:
raise ValueError("Unknown socket type.")
if plugins is None:
self.xmpp.register_plugins()
else:
for plugin in plugins:
self.xmpp.register_plugin(plugin)
# Some plugins require messages to have ID values. Set
# this to True in tests related to those plugins.
self.xmpp.use_message_ids = False
self.xmpp.process(threaded=True)
if skip:
if socket != 'live':
# Mark send queue as usable
self.xmpp.session_started_event.set()
# Clear startup stanzas
self.xmpp.socket.next_sent(timeout=1)
if mode == 'component':
self.xmpp.socket.next_sent(timeout=1)
else:
skip_queue.get(block=True, timeout=10)
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 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 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 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 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
def stream_start(self, mode='client', skip=True, header=None,
socket='mock', jid='tester@localhost',
password='******', server='localhost',
port=5222, sasl_mech=None,
plugins=None, plugin_config={}):
"""
Initialize an XMPP client or component using a dummy XML stream.
Arguments:
mode -- Either 'client' or 'component'. Defaults to 'client'.
skip -- Indicates if the first item in the sent queue (the
stream header) should be removed. Tests that wish
to test initializing the stream should set this to
False. Otherwise, the default of True should be used.
socket -- Either 'mock' or 'live' to indicate if the socket
should be a dummy, mock socket or a live, functioning
socket. Defaults to 'mock'.
jid -- The JID to use for the connection.
Defaults to 'tester@localhost'.
password -- The password to use for the connection.
Defaults to 'test'.
server -- The name of the XMPP server. Defaults to 'localhost'.
port -- The port to use when connecting to the server.
Defaults to 5222.
plugins -- List of plugins to register. By default, all plugins
are loaded.
"""
if mode == 'client':
self.xmpp = ClientXMPP(jid, password,
sasl_mech=sasl_mech,
plugin_config=plugin_config)
elif mode == 'component':
self.xmpp = ComponentXMPP(jid, password,
server, port,
plugin_config=plugin_config)
else:
raise ValueError("Unknown XMPP connection mode.")
# Remove unique ID prefix to make it easier to test
self.xmpp._id_prefix = ''
self.xmpp._disconnect_wait_for_threads = False
self.xmpp.default_lang = None
self.xmpp.peer_default_lang = None
# We will use this to wait for the session_start event
# for live connections.
skip_queue = Queue()
if socket == 'mock':
self.xmpp.set_socket(TestSocket())
# Simulate connecting for mock sockets.
self.xmpp.auto_reconnect = False
self.xmpp.state._set_state('connected')
# Must have the stream header ready for xmpp.process() to work.
if not header:
header = self.xmpp.stream_header
self.xmpp.socket.recv_data(header)
elif socket == 'live':
self.xmpp.socket_class = TestLiveSocket
def wait_for_session(x):
self.xmpp.socket.clear()
skip_queue.put('started')
self.xmpp.add_event_handler('session_start', wait_for_session)
if server is not None:
self.xmpp.connect((server, port))
else:
self.xmpp.connect()
else:
raise ValueError("Unknown socket type.")
if plugins is None:
self.xmpp.register_plugins()
else:
for plugin in plugins:
self.xmpp.register_plugin(plugin)
# Some plugins require messages to have ID values. Set
# this to True in tests related to those plugins.
self.xmpp.use_message_ids = False
self.xmpp.process(threaded=True)
if skip:
if socket != 'live':
# Mark send queue as usable
self.xmpp.session_started_event.set()
# Clear startup stanzas
self.xmpp.socket.next_sent(timeout=1)
if mode == 'component':
self.xmpp.socket.next_sent(timeout=1)
else:
skip_queue.get(block=True, timeout=10)
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 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 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