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