1#!/usr/bin/env python3
2# -*- coding: utf-8 -*-
3
4"""
5    Slixmpp: The Slick XMPP Library
6    Copyright (C) 2012  Nathanael C. Fritz
7    This file is part of Slixmpp.
8
9    See the file LICENSE for copying permission.
10"""
11
12import logging
13from getpass import getpass
14from argparse import ArgumentParser
15
16import slixmpp
17from slixmpp.exceptions import XMPPError
18from slixmpp import asyncio
19
20
21FILE_TYPES = {
22    'image/png': 'png',
23    'image/gif': 'gif',
24    'image/jpeg': 'jpg'
25}
26
27
28class AvatarDownloader(slixmpp.ClientXMPP):
29
30    """
31    A basic script for downloading the avatars for a user's contacts.
32    """
33
34    def __init__(self, jid, password):
35        slixmpp.ClientXMPP.__init__(self, jid, password)
36        self.add_event_handler("session_start", self.start)
37        self.add_event_handler("changed_status", self.wait_for_presences)
38
39        self.add_event_handler('vcard_avatar_update', self.on_vcard_avatar)
40        self.add_event_handler('avatar_metadata_publish', self.on_avatar)
41
42        self.received = set()
43        self.presences_received = asyncio.Event()
44        self.roster_received = asyncio.Event()
45
46    def roster_received_cb(self, event):
47        self.roster_received.set()
48        self.presences_received.clear()
49
50    async def start(self, event):
51        """
52        Process the session_start event.
53
54        Typical actions for the session_start event are
55        requesting the roster and broadcasting an initial
56        presence stanza.
57
58        Arguments:
59            event -- An empty dictionary. The session_start
60                     event does not provide any additional
61                     data.
62        """
63        self.send_presence()
64        self.get_roster(callback=self.roster_received_cb)
65
66        print('Waiting for presence updates...\n')
67        await self.roster_received.wait()
68        print('Roster received')
69        await self.presences_received.wait()
70        self.disconnect()
71
72    async def on_vcard_avatar(self, pres):
73        print("Received vCard avatar update from %s" % pres['from'].bare)
74        try:
75            result = await self['xep_0054'].get_vcard(pres['from'].bare, cached=True,
76                                                           timeout=5)
77        except XMPPError:
78            print("Error retrieving avatar for %s" % pres['from'])
79            return
80        avatar = result['vcard_temp']['PHOTO']
81
82        filetype = FILE_TYPES.get(avatar['TYPE'], 'png')
83        filename = 'vcard_avatar_%s_%s.%s' % (
84                pres['from'].bare,
85                pres['vcard_temp_update']['photo'],
86                filetype)
87        with open(filename, 'wb+') as img:
88            img.write(avatar['BINVAL'])
89
90    async def on_avatar(self, msg):
91        print("Received avatar update from %s" % msg['from'])
92        metadata = msg['pubsub_event']['items']['item']['avatar_metadata']
93        for info in metadata['items']:
94            if not info['url']:
95                try:
96                    result = await self['xep_0084'].retrieve_avatar(msg['from'].bare, info['id'],
97                                                                         timeout=5)
98                except XMPPError:
99                    print("Error retrieving avatar for %s" % msg['from'])
100                    return
101
102                avatar = result['pubsub']['items']['item']['avatar_data']
103
104                filetype = FILE_TYPES.get(metadata['type'], 'png')
105                filename = 'avatar_%s_%s.%s' % (msg['from'].bare, info['id'], filetype)
106                with open(filename, 'wb+') as img:
107                    img.write(avatar['value'])
108            else:
109                # We could retrieve the avatar via HTTP, etc here instead.
110                pass
111
112    def wait_for_presences(self, pres):
113        """
114        Wait to receive updates from all roster contacts.
115        """
116        self.received.add(pres['from'].bare)
117        print((len(self.received), len(self.client_roster.keys())))
118        if len(self.received) >= len(self.client_roster.keys()):
119            self.presences_received.set()
120        else:
121            self.presences_received.clear()
122
123
124if __name__ == '__main__':
125    # Setup the command line arguments.
126    parser = ArgumentParser()
127    parser.add_argument("-q","--quiet", help="set logging to ERROR",
128                        action="store_const",
129                        dest="loglevel",
130                        const=logging.ERROR,
131                        default=logging.ERROR)
132    parser.add_argument("-d","--debug", help="set logging to DEBUG",
133                        action="store_const",
134                        dest="loglevel",
135                        const=logging.DEBUG,
136                        default=logging.ERROR)
137
138    # JID and password options.
139    parser.add_argument("-j", "--jid", dest="jid",
140                        help="JID to use")
141    parser.add_argument("-p", "--password", dest="password",
142                        help="password to use")
143
144    args = parser.parse_args()
145
146    # Setup logging.
147    logging.basicConfig(level=args.loglevel,
148                        format='%(levelname)-8s %(message)s')
149
150    if args.jid is None:
151        args.jid = input("Username: ")
152    if args.password is None:
153        args.password = getpass("Password: ")
154
155    xmpp = AvatarDownloader(args.jid, args.password)
156    xmpp.register_plugin('xep_0054')
157    xmpp.register_plugin('xep_0153')
158    xmpp.register_plugin('xep_0084')
159
160    # Connect to the XMPP server and start processing XMPP stanzas.
161    xmpp.connect()
162    xmpp.process(forever=False)
163