code.delx.au - pymsnt/blob - src/avatar.py

   1 # Copyright 2005 James Bunton <james@delx.cjb.net>
   2 # Licensed for distribution under the GPL version 2, check COPYING for details
   3
   4 import utils
   5 import config
   6 from twisted.internet import reactor
   7 if(utils.checkTwisted()):
   8         from twisted.xish.domish import Element
   9 else:
  10         from tlib.domish import Element
  11 from debug import LogEvent, INFO, WARN, ERROR
  12 import jabw
  13 import config
  14 import lang
  15 import sha
  16 import base64
  17 import os
  18 import os.path
  19
  20 SPOOL_UMASK = 0077
  21
  22 def parsePhotoEl(photo):
  23         """ Pass the photo element as an avatar, returns the avatar imageData """
  24         imageData = ""
  25         imageType = ""
  26         for e in photo.elements():
  27                 if(e.name == "BINVAL"):
  28                         imageData = base64.decodestring(e.__str__())
  29                 elif(e.name == "TYPE"):
  30                         imageType = e.__str__()
  31
  32         if(imageType != "image/png"):
  33                 imageData = utils.convertToPNG(imageData)
  34
  35         return imageData
  36
  37
  38
  39 class Avatar:
  40         """ Represents an Avatar. Does not store the image in memory. """
  41         def __init__(self, imageData, avatarCache):
  42                 self.__imageHash = sha.sha(imageData).hexdigest()
  43                 self.__avatarCache = avatarCache
  44
  45         def getImageHash(self):
  46                 """ Returns the SHA1 hash of the avatar. """
  47                 return self.__imageHash
  48
  49         def getImageData(self):
  50                 """ Returns this Avatar's imageData. This loads data from a file. """
  51                 return self.__avatarCache.getAvatarData(self.__imageHash)
  52
  53         def makePhotoElement(self):
  54                 """ Returns an XML Element that can be put into the vCard. """
  55                 photo = Element((None, "PHOTO"))
  56                 type = photo.addElement("TYPE")
  57                 type.addContent("image/png")
  58                 binval = photo.addElement("BINVAL")
  59                 binval.addContent(base64.encodestring(self.getImageData()))
  60                 return photo
  61
  62         def makeDataElement(self):
  63                 """ Returns an XML Element that can be put into a jabber:x:avatar IQ stanza. """
  64                 data = Element((None, "data"))
  65                 data["mimetype"] = "image/png"
  66                 data.addContent(base64.encodestring(self.getImageData()))
  67                 return data
  68
  69         def __eq__(self, other):
  70                 return (other and self.__imageHash == other.__imageHash)
  71
  72
  73 class AvatarCache:
  74         """ Manages avatars on disk. Avatars are stored according to their SHA1 hash.
  75         The layout is config.spooldir / config.jid / avatars / "first two characters of SHA1 hash" """
  76
  77         def dir(self, key):
  78                 """ Returns the full path to the directory that a
  79                 particular key is in. Creates that directory if it doesn't already exist. """
  80                 d = os.path.abspath(config.spooldir) + "/" + config.jid + "/avatars/" + key[0:3] + "/"
  81                 if not os.path.exists(d):
  82                         os.makedirs(d)
  83                 return d
  84
  85         def setAvatar(self, imageData):
  86                 """ Writes an avatar to disk according to its key.
  87                 Returns an Avatar object. """
  88                 avatar = Avatar(imageData, self)
  89                 key = avatar.getImageHash()
  90                 LogEvent(INFO, "", "Setting avatar %s" % (key))
  91                 prev_umask = os.umask(SPOOL_UMASK)
  92                 try:
  93                         f = open(self.dir(key) + key, 'w')
  94                         f.write(imageData)
  95                         f.close()
  96                 except IOError, e:
  97                         LogEvent(WARN, "", "IOError writing to avatar %s - %s" % (key, str(e)))
  98                 os.umask(prev_umask)
  99                 return avatar
 100
 101         def getAvatar(self, key):
 102                 """ Loads the avatar with SHA1 hash of 'key' from disk and returns an Avatar object """
 103                 imageData = self.getAvatarData(key)
 104                 if imageData:
 105                         return Avatar(imageData, self)
 106
 107         def getAvatarData(self, key):
 108                 """ Loads the avatar with SHA1 hash of 'key' from disk and returns the data """
 109                 try:
 110                         filename = self.dir(key) + key
 111                         if os.path.isfile(filename):
 112                                 LogEvent(INFO, "Getting avatar.")
 113                                 f = open(filename)
 114                                 data = f.read()
 115                                 f.close()
 116                                 return data
 117                         else:
 118                                 LogEvent(INFO, "", "Avatar not found.")
 119                 except IOError, e:
 120                         LogEvent(WARN, "", "IOError reading avatar.")
 121
 122
 123