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