roster.py 8.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184
  1. ## roster.py
  2. ##
  3. ## Copyright (C) 2003-2005 Alexey "Snake" Nezhdanov
  4. ##
  5. ## This program is free software; you can redistribute it and/or modify
  6. ## it under the terms of the GNU General Public License as published by
  7. ## the Free Software Foundation; either version 2, or (at your option)
  8. ## 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. # $Id$
  15. """
  16. Simple roster implementation. Can be used though for different tasks like
  17. mass-renaming of contacts.
  18. """
  19. from .protocol import *
  20. from .client import PlugIn
  21. class Roster(PlugIn):
  22. """ Defines a plenty of methods that will allow you to manage roster.
  23. Also automatically track presences from remote JIDs taking into
  24. account that every JID can have multiple resources connected. Does not
  25. currently support 'error' presences.
  26. You can also use mapping interface for access to the internal representation of
  27. contacts in roster.
  28. """
  29. def __init__(self):
  30. """ Init internal variables. """
  31. PlugIn.__init__(self)
  32. self.DBG_LINE='roster'
  33. self._data = {}
  34. self.set=None
  35. self._exported_methods=[self.getRoster]
  36. def plugin(self,owner,request=1):
  37. """ Register presence and subscription trackers in the owner's dispatcher.
  38. Also request roster from server if the 'request' argument is set.
  39. Used internally."""
  40. self._owner.RegisterHandler('iq',self.RosterIqHandler,'result',NS_ROSTER)
  41. self._owner.RegisterHandler('iq',self.RosterIqHandler,'set',NS_ROSTER)
  42. self._owner.RegisterHandler('presence',self.PresenceHandler)
  43. if request: self.Request()
  44. def Request(self,force=0):
  45. """ Request roster from server if it were not yet requested
  46. (or if the 'force' argument is set). """
  47. if self.set is None: self.set=0
  48. elif not force: return
  49. self._owner.send(Iq('get',NS_ROSTER))
  50. self.DEBUG('Roster requested from server','start')
  51. def getRoster(self):
  52. """ Requests roster from server if neccessary and returns self."""
  53. if not self.set: self.Request()
  54. while not self.set: self._owner.Process(10)
  55. return self
  56. def RosterIqHandler(self,dis,stanza):
  57. """ Subscription tracker. Used internally for setting items state in
  58. internal roster representation. """
  59. for item in stanza.getTag('query').getTags('item'):
  60. jid=item.getAttr('jid')
  61. if item.getAttr('subscription')=='remove':
  62. if jid in self._data: del self._data[jid]
  63. raise NodeProcessed # a MUST
  64. self.DEBUG('Setting roster item %s...'%jid,'ok')
  65. if jid not in self._data: self._data[jid]={}
  66. self._data[jid]['name']=item.getAttr('name')
  67. self._data[jid]['ask']=item.getAttr('ask')
  68. self._data[jid]['subscription']=item.getAttr('subscription')
  69. self._data[jid]['groups']=[]
  70. if 'resources' not in self._data[jid]: self._data[jid]['resources']={}
  71. for group in item.getTags('group'): self._data[jid]['groups'].append(group.getData())
  72. self._data[self._owner.User+'@'+self._owner.Server]={'resources':{},'name':None,'ask':None,'subscription':None,'groups':None,}
  73. self.set=1
  74. raise NodeProcessed # a MUST. Otherwise you'll get back an <iq type='error'/>
  75. def PresenceHandler(self,dis,pres):
  76. """ Presence tracker. Used internally for setting items' resources state in
  77. internal roster representation. """
  78. jid=JID(pres.getFrom())
  79. if jid.getStripped() not in self._data: self._data[jid.getStripped()]={'name':None,'ask':None,'subscription':'none','groups':['Not in roster'],'resources':{}}
  80. item=self._data[jid.getStripped()]
  81. typ=pres.getType()
  82. if not typ:
  83. self.DEBUG('Setting roster item %s for resource %s...'%(jid.getStripped(),jid.getResource()),'ok')
  84. item['resources'][jid.getResource()]=res={'show':None,'status':None,'priority':'0','timestamp':None}
  85. if pres.getTag('show'): res['show']=pres.getShow()
  86. if pres.getTag('status'): res['status']=pres.getStatus()
  87. if pres.getTag('priority'): res['priority']=pres.getPriority()
  88. if not pres.getTimestamp(): pres.setTimestamp()
  89. res['timestamp']=pres.getTimestamp()
  90. elif typ=='unavailable' and jid.getResource() in item['resources']: del item['resources'][jid.getResource()]
  91. # Need to handle type='error' also
  92. def _getItemData(self,jid,dataname):
  93. """ Return specific jid's representation in internal format. Used internally. """
  94. jid=jid[:(jid+'/').find('/')]
  95. return self._data[jid][dataname]
  96. def _getResourceData(self,jid,dataname):
  97. """ Return specific jid's resource representation in internal format. Used internally. """
  98. if jid.find('/')+1:
  99. jid,resource=jid.split('/',1)
  100. if resource in self._data[jid]['resources']: return self._data[jid]['resources'][resource][dataname]
  101. elif list(self._data[jid]['resources'].keys()):
  102. lastpri=-129
  103. for r in list(self._data[jid]['resources'].keys()):
  104. if int(self._data[jid]['resources'][r]['priority'])>lastpri: resource,lastpri=r,int(self._data[jid]['resources'][r]['priority'])
  105. return self._data[jid]['resources'][resource][dataname]
  106. def delItem(self,jid):
  107. """ Delete contact 'jid' from roster."""
  108. self._owner.send(Iq('set',NS_ROSTER,payload=[Node('item',{'jid':jid,'subscription':'remove'})]))
  109. def getAsk(self,jid):
  110. """ Returns 'ask' value of contact 'jid'."""
  111. return self._getItemData(jid,'ask')
  112. def getGroups(self,jid):
  113. """ Returns groups list that contact 'jid' belongs to."""
  114. return self._getItemData(jid,'groups')
  115. def getName(self,jid):
  116. """ Returns name of contact 'jid'."""
  117. return self._getItemData(jid,'name')
  118. def getPriority(self,jid):
  119. """ Returns priority of contact 'jid'. 'jid' should be a full (not bare) JID."""
  120. return self._getResourceData(jid,'priority')
  121. def getRawRoster(self):
  122. """ Returns roster representation in internal format. """
  123. return self._data
  124. def getRawItem(self,jid):
  125. """ Returns roster item 'jid' representation in internal format. """
  126. return self._data[jid[:(jid+'/').find('/')]]
  127. def getShow(self, jid):
  128. """ Returns 'show' value of contact 'jid'. 'jid' should be a full (not bare) JID."""
  129. return self._getResourceData(jid,'show')
  130. def getStatus(self, jid):
  131. """ Returns 'status' value of contact 'jid'. 'jid' should be a full (not bare) JID."""
  132. return self._getResourceData(jid,'status')
  133. def getSubscription(self,jid):
  134. """ Returns 'subscription' value of contact 'jid'."""
  135. return self._getItemData(jid,'subscription')
  136. def getResources(self,jid):
  137. """ Returns list of connected resources of contact 'jid'."""
  138. return list(self._data[jid[:(jid+'/').find('/')]]['resources'].keys())
  139. def setItem(self,jid,name=None,groups=[]):
  140. """ Creates/renames contact 'jid' and sets the groups list that it now belongs to."""
  141. iq=Iq('set',NS_ROSTER)
  142. query=iq.getTag('query')
  143. attrs={'jid':jid}
  144. if name: attrs['name']=name
  145. item=query.setTag('item',attrs)
  146. for group in groups: item.addChild(node=Node('group',payload=[group]))
  147. self._owner.send(iq)
  148. def getItems(self):
  149. """ Return list of all [bare] JIDs that the roster is currently tracks."""
  150. return list(self._data.keys())
  151. def keys(self):
  152. """ Same as getItems. Provided for the sake of dictionary interface."""
  153. return list(self._data.keys())
  154. def __getitem__(self,item):
  155. """ Get the contact in the internal format. Raises KeyError if JID 'item' is not in roster."""
  156. return self._data[item]
  157. def getItem(self,item):
  158. """ Get the contact in the internal format (or None if JID 'item' is not in roster)."""
  159. if item in self._data: return self._data[item]
  160. def Subscribe(self,jid):
  161. """ Send subscription request to JID 'jid'."""
  162. self._owner.send(Presence(jid,'subscribe'))
  163. def Unsubscribe(self,jid):
  164. """ Ask for removing our subscription for JID 'jid'."""
  165. self._owner.send(Presence(jid,'unsubscribe'))
  166. def Authorize(self,jid):
  167. """ Authorise JID 'jid'. Works only if these JID requested auth previously. """
  168. self._owner.send(Presence(jid,'subscribed'))
  169. def Unauthorize(self,jid):
  170. """ Unauthorise JID 'jid'. Use for declining authorisation request
  171. or for removing existing authorization. """
  172. self._owner.send(Presence(jid,'unsubscribed'))