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