net/tools/testserver/xmppserver.py

#!/usr/bin/python2.4
# Copyright (c) 2010 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""A bare-bones and non-compliant XMPP server.

Just enough of the protocol is implemented to get it to work with
Chrome's sync notification system.
"""

import asynchat
import asyncore
import base64
import re
import socket
from xml.dom import minidom

# pychecker complains about the use of fileno(), which is implemented
# by asyncore by forwarding to an internal object via __getattr__.
__pychecker__ = 'no-classattr'


class Error(Exception):
  """Error class for this module."""
  pass


class UnexpectedXml(Error):
  """Raised when an unexpected XML element has been encountered."""

  def __init__(self, xml_element):
    xml_text = xml_element.toxml()
    Error.__init__(self, 'Unexpected XML element', xml_text)


def ParseXml(xml_string):
  """Parses the given string as XML and returns a minidom element
  object.
  """
  dom = minidom.parseString(xml_string)

  # minidom handles xmlns specially, but there's a bug where it sets
  # the attribute value to None, which causes toxml() or toprettyxml()
  # to break.
  def FixMinidomXmlnsBug(xml_element):
    if xml_element.getAttribute('xmlns') is None:
      xml_element.setAttribute('xmlns', '')

  def ApplyToAllDescendantElements(xml_element, fn):
    fn(xml_element)
    for node in xml_element.childNodes:
      if node.nodeType == node.ELEMENT_NODE:
        ApplyToAllDescendantElements(node, fn)

  root = dom.documentElement
  ApplyToAllDescendantElements(root, FixMinidomXmlnsBug)
  return root


def CloneXml(xml):
  """Returns a deep copy of the given XML element.

  Args:
    xml: The XML element, which should be something returned from
         ParseXml() (i.e., a root element).
  """
  return xml.ownerDocument.cloneNode(True).documentElement


class StanzaParser(object):
  """A hacky incremental XML parser.

  StanzaParser consumes data incrementally via FeedString() and feeds
  its delegate complete parsed stanzas (i.e., XML documents) via
  FeedStanza().  Any stanzas passed to FeedStanza() are unlinked after
  the callback is done.

  Use like so:

  class MyClass(object):
    ...
    def __init__(self, ...):
      ...
      self._parser = StanzaParser(self)
      ...

    def SomeFunction(self, ...):
      ...
      self._parser.FeedString(some_data)
      ...

    def FeedStanza(self, stanza):
      ...
      print stanza.toprettyxml()
      ...
  """

  # NOTE(akalin): The following regexps are naive, but necessary since
  # none of the existing Python 2.4/2.5 XML libraries support
  # incremental parsing.  This works well enough for our purposes.
  #
  # The regexps below assume that any present XML element starts at
  # the beginning of the string, but there may be trailing whitespace.

  # Matches an opening stream tag (e.g., '<stream:stream foo="bar">')
  # (assumes that the stream XML namespace is defined in the tag).
  _stream_re = re.compile(r'^(<stream:stream [^>]*>)\s*')

  # Matches an empty element tag (e.g., '<foo bar="baz"/>').
  _empty_element_re = re.compile(r'^(<[^>]*/>)\s*')

  # Matches a non-empty element (e.g., '<foo bar="baz">quux</foo>').
  # Does *not* handle nested elements.
  _non_empty_element_re = re.compile(r'^(<([^ >]*)[^>]*>.*?</\2>)\s*')

  # The closing tag for a stream tag.  We have to insert this
  # ourselves since all XML stanzas are children of the stream tag,
  # which is never closed until the connection is closed.
  _stream_suffix = '</stream:stream>'

  def __init__(self, delegate):
    self._buffer = ''
    self._delegate = delegate

  def FeedString(self, data):
    """Consumes the given string data, possibly feeding one or more
    stanzas to the delegate.
    """
    self._buffer += data
    while (self._ProcessBuffer(self._stream_re, self._stream_suffix) or
           self._ProcessBuffer(self._empty_element_re) or
           self._ProcessBuffer(self._non_empty_element_re)):
      pass

  def _ProcessBuffer(self, regexp, xml_suffix=''):
    """If the buffer matches the given regexp, removes the match from
    the buffer, appends the given suffix, parses it, and feeds it to
    the delegate.

    Returns:
      Whether or not the buffer matched the given regexp.
    """
    results = regexp.match(self._buffer)
    if not results:
      return False
    xml_text = self._buffer[:results.end()] + xml_suffix
    self._buffer = self._buffer[results.end():]
    stanza = ParseXml(xml_text)
    self._delegate.FeedStanza(stanza)
    # Needed because stanza may have cycles.
    stanza.unlink()
    return True


class Jid(object):
  """Simple struct for an XMPP jid (essentially an e-mail address with
  an optional resource string).
  """

  def __init__(self, username, domain, resource=''):
    self.username = username
    self.domain = domain
    self.resource = resource

  def __str__(self):
    jid_str = "%s@%s" % (self.username, self.domain)
    if self.resource:
      jid_str += '/' + self.resource
    return jid_str

  def GetBareJid(self):
    return Jid(self.username, self.domain)


class IdGenerator(object):
  """Simple class to generate unique IDs for XMPP messages."""

  def __init__(self, prefix):
    self._prefix = prefix
    self._id = 0

  def GetNextId(self):
    next_id = "%s.%s" % (self._prefix, self._id)
    self._id += 1
    return next_id


class HandshakeTask(object):
  """Class to handle the initial handshake with a connected XMPP
  client.
  """

  # The handshake states in order.
  (_INITIAL_STREAM_NEEDED,
   _AUTH_NEEDED,
   _AUTH_STREAM_NEEDED,
   _BIND_NEEDED,
   _SESSION_NEEDED,
   _FINISHED) = range(6)

  # Used when in the _INITIAL_STREAM_NEEDED and _AUTH_STREAM_NEEDED
  # states.  Not an XML object as it's only the opening tag.
  #
  # The from and id attributes are filled in later.
  _STREAM_DATA = (
    '<stream:stream from="%s" id="%s" '
    'version="1.0" xmlns:stream="http://etherx.jabber.org/streams" '
    'xmlns="jabber:client">')

  # Used when in the _INITIAL_STREAM_NEEDED state.
  _AUTH_STANZA = ParseXml(
    '<stream:features xmlns:stream="http://etherx.jabber.org/streams">'
    '  <mechanisms xmlns="urn:ietf:params:xml:ns:xmpp-sasl">'
    '    <mechanism>PLAIN</mechanism>'
    '    <mechanism>X-GOOGLE-TOKEN</mechanism>'
    '  </mechanisms>'
    '</stream:features>')

  # Used when in the _AUTH_NEEDED state.
  _AUTH_SUCCESS_STANZA = ParseXml(
    '<success xmlns="urn:ietf:params:xml:ns:xmpp-sasl"/>')

  # Used when in the _AUTH_STREAM_NEEDED state.
  _BIND_STANZA = ParseXml(
    '<stream:features xmlns:stream="http://etherx.jabber.org/streams">'
    '  <bind xmlns="urn:ietf:params:xml:ns:xmpp-bind"/>'
    '  <session xmlns="urn:ietf:params:xml:ns:xmpp-session"/>'
    '</stream:features>')

  # Used when in the _BIND_NEEDED state.
  #
  # The id and jid attributes are filled in later.
  _BIND_RESULT_STANZA = ParseXml(
    '<iq id="" type="result">'
    '  <bind xmlns="urn:ietf:params:xml:ns:xmpp-bind">'
    '    <jid/>'
    '  </bind>'
    '</iq>')

  # Used when in the _SESSION_NEEDED state.
  #
  # The id attribute is filled in later.
  _IQ_RESPONSE_STANZA = ParseXml('<iq id="" type="result"/>')

  def __init__(self, connection, id_generator, resource_prefix):
    self._connection = connection
    self._id_generator = id_generator
    self._username = ''
    self._domain = ''
    self._jid = None
    self._resource_prefix = resource_prefix
    self._state = self._INITIAL_STREAM_NEEDED

  def FeedStanza(self, stanza):
    """Inspects the given stanza and changes the handshake state if needed.

    Called when a stanza is received from the client.  Inspects the
    stanza to make sure it has the expected attributes given the
    current state, advances the state if needed, and sends a reply to
    the client if needed.
    """
    def ExpectStanza(stanza, name):
      if stanza.tagName != name:
        raise UnexpectedXml(stanza)

    def ExpectIq(stanza, type, name):
      ExpectStanza(stanza, 'iq')
      if (stanza.getAttribute('type') != type or
          stanza.firstChild.tagName != name):
        raise UnexpectedXml(stanza)

    def GetStanzaId(stanza):
      return stanza.getAttribute('id')

    def HandleStream(stanza):
      ExpectStanza(stanza, 'stream:stream')
      domain = stanza.getAttribute('to')
      if domain:
        self._domain = domain
      SendStreamData()

    def SendStreamData():
      next_id = self._id_generator.GetNextId()
      stream_data = self._STREAM_DATA % (self._domain, next_id)
      self._connection.SendData(stream_data)

    def GetUserDomain(stanza):
      encoded_username_password = stanza.firstChild.data
      username_password = base64.b64decode(encoded_username_password)
      (_, username_domain, _) = username_password.split('\0')
      # The domain may be omitted.
      #
      # If we were using python 2.5, we'd be able to do:
      #
      #   username, _, domain = username_domain.partition('@')
      #   if not domain:
      #     domain = self._domain
      at_pos = username_domain.find('@')
      if at_pos != -1:
        username = username_domain[:at_pos]
        domain = username_domain[at_pos+1:]
      else:
        username = username_domain
        domain = self._domain
      return (username, domain)

    if self._state == self._INITIAL_STREAM_NEEDED:
      HandleStream(stanza)
      self._connection.SendStanza(self._AUTH_STANZA, False)
      self._state = self._AUTH_NEEDED

    elif self._state == self._AUTH_NEEDED:
      ExpectStanza(stanza, 'auth')
      (self._username, self._domain) = GetUserDomain(stanza)
      self._connection.SendStanza(self._AUTH_SUCCESS_STANZA, False)
      self._state = self._AUTH_STREAM_NEEDED

    elif self._state == self._AUTH_STREAM_NEEDED:
      HandleStream(stanza)
      self._connection.SendStanza(self._BIND_STANZA, False)
      self._state = self._BIND_NEEDED

    elif self._state == self._BIND_NEEDED:
      ExpectIq(stanza, 'set', 'bind')
      stanza_id = GetStanzaId(stanza)
      resource_element = stanza.getElementsByTagName('resource')[0]
      resource = resource_element.firstChild.data
      full_resource = '%s.%s' % (self._resource_prefix, resource)
      response = CloneXml(self._BIND_RESULT_STANZA)
      response.setAttribute('id', stanza_id)
      self._jid = Jid(self._username, self._domain, full_resource)
      jid_text = response.parentNode.createTextNode(str(self._jid))
      response.getElementsByTagName('jid')[0].appendChild(jid_text)
      self._connection.SendStanza(response)
      self._state = self._SESSION_NEEDED

    elif self._state == self._SESSION_NEEDED:
      ExpectIq(stanza, 'set', 'session')
      stanza_id = GetStanzaId(stanza)
      xml = CloneXml(self._IQ_RESPONSE_STANZA)
      xml.setAttribute('id', stanza_id)
      self._connection.SendStanza(xml)
      self._state = self._FINISHED
      self._connection.HandshakeDone(self._jid)


def AddrString(addr):
  return '%s:%d' % addr


class XmppConnection(asynchat.async_chat):
  """A single XMPP client connection.

  This class handles the connection to a single XMPP client (via a
  socket).  It does the XMPP handshake and also implements the (old)
  Google notification protocol.
  """

  # We use this XML template for subscription responses as well as
  # notifications (conveniently enough, the same template works
  # for both).
  #
  # The from, to, id, and type attributes are filled in later.
  _NOTIFIER_STANZA = ParseXml(
    """<iq from="" to="" id="" type="">
         <not:getAll xmlns:not="google:notifier">
           <Result xmlns=""/>
         </not:getAll>
       </iq>
    """)

  def __init__(self, sock, socket_map, connections, addr):
    """Starts up the xmpp connection.

    Args:
      sock: The socket to the client.
      socket_map: A map from sockets to their owning objects.
      connections: The set of handshake-completed connections.
      addr: The host/port of the client.
    """
    asynchat.async_chat.__init__(self, sock)
    self.set_terminator(None)
    # async_chat in Python 2.4 has a bug where it ignores a
    # socket_map argument.  So we handle that ourselves.
    self._socket_map = socket_map
    self._socket_map[self.fileno()] = self

    self._connections = connections
    self._parser = StanzaParser(self)
    self._jid = None

    self._addr = addr
    addr_str = AddrString(self._addr)
    self._id_generator = IdGenerator(addr_str)
    self._handshake_task = (
      HandshakeTask(self, self._id_generator, addr_str))
    print 'Starting connection to %s' % self

  def __str__(self):
    if self._jid:
      return str(self._jid)
    else:
      return AddrString(self._addr)

  # async_chat implementation.

  def collect_incoming_data(self, data):
    self._parser.FeedString(data)

  # This is only here to make pychecker happy.
  def found_terminator(self):
    asynchat.async_chat.found_terminator(self)

  def handle_close(self):
    print "Closing connection to %s" % self
    # Remove ourselves from anywhere we possibly installed ourselves.
    self._connections.discard(self)
    del self._socket_map[self.fileno()]

  # Called by self._parser.FeedString().
  def FeedStanza(self, stanza):
    if self._handshake_task:
      self._handshake_task.FeedStanza(stanza)
    elif stanza.tagName == 'iq':
      self._HandleIq(stanza)
    else:
      raise UnexpectedXml(stanza)

  # Called by self._handshake_task.
  def HandshakeDone(self, jid):
    self._jid = jid
    self._handshake_task = None
    self._connections.add(self)
    print "Handshake done for %s" % self

  def _HandleIq(self, iq):
    if (iq.firstChild and
        iq.firstChild.namespaceURI == 'google:notifier'):
      iq_id = iq.getAttribute('id')
      self._HandleNotifierCommand(iq_id, iq.firstChild)
    elif iq.getAttribute('type') == 'result':
      # Ignore all client acks.
      pass
    else:
      raise UnexpectedXml(iq)

  def _HandleNotifierCommand(self, id, command_xml):
    command = command_xml.tagName
    if command == 'getAll':
      # Subscription request.
      if not command_xml.getElementsByTagName('SubscribedServiceUrl'):
        raise UnexpectedXml(command_xml)
      self._SendNotifierStanza(id, 'result')
    elif command == 'set':
      # Send notification request.
      SendNotification(self._connections)
    else:
      raise UnexpectedXml(command_xml)

  def _SendNotifierStanza(self, id, type):
    stanza = CloneXml(self._NOTIFIER_STANZA)
    stanza.setAttribute('from', str(self._jid.GetBareJid()))
    stanza.setAttribute('to', str(self._jid))
    stanza.setAttribute('id', id)
    stanza.setAttribute('type', type)
    self.SendStanza(stanza)

  def SendStanza(self, stanza, unlink=True):
    """Sends a stanza to the client.

    Args:
      stanza: The stanza to send.
      unlink: Whether to unlink stanza after sending it. (Pass in
      False if stanza is a constant.)
    """
    self.SendData(stanza.toxml())
    if unlink:
      stanza.unlink()

  def SendData(self, data):
    """Sends raw data to the client.
    """
    # We explicitly encode to ascii as that is what the client expects
    # (some minidom library functions return unicode strings).
    self.push(data.encode('ascii'))

  def SendNotification(self):
    """Sends a notification to the client."""
    next_id = self._id_generator.GetNextId()
    self._SendNotifierStanza(next_id, 'set')


def SendNotification(connections):
  """Sends a notification to all connections in the given sequence."""
  for connection in connections:
    print 'Sending notification to %s' % connection
    connection.SendNotification()


class XmppServer(asyncore.dispatcher):
  """The main XMPP server class.

  The XMPP server starts accepting connections on the given address
  and spawns off XmppConnection objects for each one.

  Use like so:

    socket_map = {}
    xmpp_server = xmppserver.XmppServer(socket_map, ('127.0.0.1', 5222))
    asyncore.loop(30.0, False, socket_map)
  """

  def __init__(self, socket_map, addr):
    asyncore.dispatcher.__init__(self, None, socket_map)
    self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
    self.set_reuse_addr()
    self.bind(addr)
    self.listen(5)
    self._socket_map = socket_map
    self._socket_map[self.fileno()] = self
    self._connections = set()

  def handle_accept(self):
    (sock, addr) = self.accept()
    XmppConnection(sock, self._socket_map, self._connections, addr)