src/python/osrf/net.py

   1 # -----------------------------------------------------------------------
   2 # Copyright (C) 2007  Georgia Public Library Service
   3 # Bill Erickson <billserickson@gmail.com>
   4 #
   5 # This program is free software; you can redistribute it and/or
   6 # modify it under the terms of the GNU General Public License
   7 # as published by the Free Software Foundation; either version 2
   8 # of the License, or (at your option) any later version.
   9 #
  10 # This program is distributed in the hope that it will be useful,
  11 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  12 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  13 # GNU General Public License for more details.
  14 # -----------------------------------------------------------------------
  15
  16
  17 import os, time, threading
  18 from pyxmpp.jabber.client import JabberClient
  19 from pyxmpp.message import Message
  20 from pyxmpp.jid import JID
  21 from socket import gethostname
  22 import libxml2
  23 import osrf.log
  24
  25 THREAD_SESSIONS = {}
  26
  27 # - log jabber activity (for future reference)
  28 #import logging
  29 #logger=logging.getLogger()
  30 #logger.addHandler(logging.StreamHandler())
  31 #logger.addHandler(logging.FileHandler('j.log'))
  32 #logger.setLevel(logging.DEBUG)
  33
  34 def set_network_handle(handle):
  35     """ Sets the thread-specific network handle"""
  36     THREAD_SESSIONS[threading.currentThread().getName()] = handle
  37
  38 def get_network_handle():
  39     """ Returns the thread-specific network connection handle."""
  40     return THREAD_SESSIONS.get(threading.currentThread().getName())
  41
  42 def clear_network_handle():
  43     ''' Disconnects the thread-specific handle and discards it '''
  44     handle = THREAD_SESSIONS.get(threading.currentThread().getName())
  45     if handle:
  46         handle.disconnect()
  47         del THREAD_SESSIONS[threading.currentThread().getName()]
  48
  49 class NetworkMessage(object):
  50     """Network message
  51
  52     attributes:
  53
  54     sender - message sender
  55     recipient - message recipient
  56     body - the body of the message
  57     thread - the message thread
  58     locale - locale of the message
  59     """
  60
  61     def __init__(self, message=None, **args):
  62         if message:
  63             self.body = message.get_body()
  64             self.thread = message.get_thread()
  65             self.recipient = message.get_to()
  66             self.router_command = None
  67             self.router_class = None
  68             if message.xmlnode.hasProp('router_from') and \
  69                 message.xmlnode.prop('router_from') != '':
  70                 self.sender = message.xmlnode.prop('router_from')
  71             else:
  72                 self.sender = message.get_from().as_utf8()
  73         else:
  74             self.sender = args.get('sender')
  75             self.recipient = args.get('recipient')
  76             self.body = args.get('body')
  77             self.thread = args.get('thread')
  78             self.router_command = args.get('router_command')
  79             self.router_class = args.get('router_class')
  80
  81     @staticmethod
  82     def from_xml(xml):
  83         doc = libxml2.parseDoc(xml)
  84         msg = Message(doc.getRootElement())
  85         return NetworkMessage(msg)
  86
  87
  88     def make_xmpp_msg(self):
  89         ''' Creates a pyxmpp.message.Message and adds custom attributes '''
  90
  91         msg = Message(None, self.sender, self.recipient, None, None, None, \
  92             self.body, self.thread)
  93         if self.router_command:
  94             msg.xmlnode.newProp('router_command', self.router_command)
  95         if self.router_class:
  96             msg.xmlnode.newProp('router_class', self.router_class)
  97         return msg
  98
  99     def to_xml(self):
 100         ''' Turns this message into XML '''
 101         return self.make_xmpp_msg().serialize()
 102
 103
 104 class Network(JabberClient):
 105     def __init__(self, **args):
 106         self.isconnected = False
 107
 108         # Create a unique jabber resource
 109         resource = args.get('resource') or 'python_client'
 110         resource += '_' + gethostname() + ':' + str(os.getpid()) + '_' + \
 111             threading.currentThread().getName().lower()
 112         self.jid = JID(args['username'], args['host'], resource)
 113
 114         osrf.log.log_debug("initializing network with JID %s and host=%s, "
 115             "port=%s, username=%s" % (self.jid.as_utf8(), args['host'], \
 116             args['port'], args['username']))
 117
 118         #initialize the superclass
 119         JabberClient.__init__(self, self.jid, args['password'], args['host'])
 120         self.queue = []
 121
 122         self.receive_callback = None
 123
 124     def connect(self):
 125         JabberClient.connect(self)
 126         while not self.isconnected:
 127             stream = self.get_stream()
 128             act = stream.loop_iter(10)
 129             if not act:
 130                 self.idle()
 131
 132     def set_receive_callback(self, func):
 133         """The callback provided is called when a message is received.
 134
 135             The only argument to the function is the received message. """
 136         self.receive_callback = func
 137
 138     def session_started(self):
 139         osrf.log.log_info("Successfully connected to the opensrf network")
 140         self.authenticated()
 141         self.stream.set_message_handler("normal", self.message_received)
 142         self.isconnected = True
 143
 144     def send(self, message):
 145         """Sends the provided network message."""
 146         osrf.log.log_internal("jabber sending to %s: %s" % (message.recipient, message.body))
 147         message.sender = self.jid.as_utf8()
 148         msg = message.make_xmpp_msg()
 149         self.stream.send(msg)
 150
 151     def message_received(self, stanza):
 152         """Handler for received messages."""
 153         if stanza.get_type()=="headline":
 154             return True
 155         # check for errors
 156         osrf.log.log_internal("jabber received message from %s : %s"
 157             % (stanza.get_from().as_utf8(), stanza.get_body()))
 158         self.queue.append(NetworkMessage(stanza))
 159         return True
 160
 161     def recv(self, timeout=120):
 162         """Attempts to receive a message from the network.
 163
 164         timeout - max number of seconds to wait for a message.
 165         If a message is received in 'timeout' seconds, the message is passed to
 166         the receive_callback is called and True is returned.  Otherwise, false is
 167         returned.
 168         """
 169
 170         forever = False
 171         if timeout < 0:
 172             forever = True
 173             timeout = None
 174
 175         if len(self.queue) == 0:
 176             while (forever or timeout >= 0) and len(self.queue) == 0:
 177                 starttime = time.time()
 178                 act = self.get_stream().loop_iter(timeout)
 179                 endtime = time.time() - starttime
 180                 if not forever:
 181                     timeout -= endtime
 182                 osrf.log.log_internal("exiting stream loop after %s seconds. "
 183                     "act=%s, queue size=%d" % (str(endtime), act, len(self.queue)))
 184                 if not act:
 185                     self.idle()
 186
 187         # if we've acquired a message, handle it
 188         msg = None
 189         if len(self.queue) > 0:
 190             msg = self.queue.pop(0)
 191             if self.receive_callback:
 192                 self.receive_callback(msg)
 193
 194         return msg
 195
 196
 197