better feedback
[samatjain:statusnet-backup.git] / StatusNet-Backup.py
1 #!/usr/bin/env python
2 # -*- coding: utf-8 -*-
3 from __future__ import print_function
4 from __future__ import unicode_literals
5
6 import argparse
7 import json
8 import os
9 import sys
10 import time
11 import urllib
12
13 import dateutil.parser
14 import requests
15
16 from lxml import etree
17
18
19 class StatusNet(object):
20     timeline_url = None
21
22     headers = {'user-agent': 'StatusNet 1.0 Backup'}
23     stream_types = ('friends_timeline', 'user_timeline', 'favorites', 'memberships', 'subscriptions')
24
25     namespaces = {
26         'activity': 'http://activitystrea.ms/spec/1.0/',
27         'app': 'http://www.w3.org/2007/app',
28         'atom': 'http://www.w3.org/2005/Atom'
29     }
30
31     urls = {} # Cache URLs for streams
32
33     def __init__(self, user_name, endpoint="http://identi.ca"):
34         self.user_name = user_name
35         self.endpoint = endpoint
36
37         user_agent = '%s (User Contact: %s/%s)' \
38             % (self.headers['user-agent'], endpoint, user_name)
39         self.headers['user-agent'] = user_agent
40
41         self.rs = requests.session(headers=self.headers, config={'max_retries': 2})
42
43     def _cacheTimelineUrls(self):
44         # Request service document
45         service_doc_url = '%s/api/statusnet/app/service/%s.xml' % (self.endpoint, self.user_name)
46         response = self.rs.get(service_doc_url)
47
48         if response is None:
49             return None
50
51         document = etree.fromstring(response.content)
52         
53         streams = document.xpath('//app:collection[@href]', namespaces=self.namespaces)
54         if not streams:
55             return None
56
57         # TODO: figure out a smarter way to write this
58         urls = {}
59         for s in streams:
60             url = s.get('href')
61             for t in self.stream_types:
62                 if t in url:
63                     urls[t] = url
64                 continue
65
66         # HACK: Don't know how to discover this—hard-coded for now
67         url = '%s/api/statuses/friends_timeline/%s.atom' % (self.endpoint, self.user_name)
68         urls['friends_timeline'] = url
69
70         self.urls = urls
71         return urls
72
73     def getTimelineUrl(self, timeline):
74         if timeline in self.urls:
75             return self.urls[timeline]
76
77         urls = self._cacheTimelineUrls()
78
79         if timeline in urls:
80             return urls[timeline]
81
82     def fetch(self, timeline, pageNo, format='atom'):
83         url = self.getTimelineUrl(timeline) + '?page=%d' % pageNo
84
85         if format in ('as', 'json'):
86             url = url.replace('.atom', '.as')
87
88         print('Getting %s' % url)
89         response = self.rs.get(url)
90
91         if response.status_code != requests.codes.ok:
92             return None
93
94         return response.content
95
96
97 def main():
98     parser = argparse.ArgumentParser(description='Backup your Identi.ca account')
99     parser.add_argument('--username', required=True)
100     parser.add_argument('--endpoint', default='http://identi.ca', help='(default: %(default)s)')
101     parser.add_argument('--timeline', choices=StatusNet.stream_types, default='user_timeline')
102     parser.add_argument('--page', type=int, default=1, help='Page number from which to start backup')
103     parser.add_argument('--force', action='store_true', default=False, help='Force overwrite, ignoring previously backed-up entries')
104
105     config = parser.parse_args()
106
107     sn = StatusNet(config.username, config.endpoint)
108     format = 'atom' # Hard-coded for now
109     skippedEntries = 0
110
111     # Create output directory
112     try:
113         os.makedirs(config.timeline)
114     except Exception:
115         pass
116     os.chdir(config.timeline)
117
118     for pageNo in range(config.page, 500):
119         print('Processing page %d' % pageNo)
120         raw_document = sn.fetch(config.timeline, pageNo, format=format)
121
122         if raw_document is None:
123             exit('Error loading page %d' % pageNo)
124         
125         if format == 'atom':
126             # Fix for: ValueError: Unicode strings with encoding declaration are not supported
127             # Dear lxml and Python: go die in a fire w/ your Unicode idiocy, thanks
128             try:
129                 raw_document = raw_document.encode('utf-8')
130             except UnicodeDecodeError:
131                 pass
132             document = etree.fromstring(raw_document)
133         elif format == 'as':
134             document = json.loads(raw_document)
135
136     #    for entry in json_document['items']:
137     #        entry_id = entry['id']
138     #        entry_text = entry['title']
139     #        print entry_id, entry_text
140
141         for entry in document.xpath('//atom:entry', namespaces=sn.namespaces):
142             entry_id = entry.xpath('atom:id', namespaces=sn.namespaces)[0].text
143
144             # Get published time & calculate file timestamp
145             published_time = entry.xpath('atom:published', namespaces=sn.namespaces)[0].text
146             published_time = dateutil.parser.parse(published_time)
147             f_time = time.mktime(published_time.timetuple())
148
149             filename = urllib.quote_plus(entry_id) + '.atom'
150
151             # We've seen this entry before
152             if os.path.isfile(filename) and not config.force:
153                 print('Skipping %s' % filename, file=sys.stderr)
154                 # this should be return to stop, continue to skip current entry
155                 # return
156                 if skippedEntries > 16:
157                     return
158                 skippedEntries = skippedEntries + 1
159                 continue
160
161             print('Backing up %s' % filename)
162             f = file(filename, 'w')
163             f.write(etree.tostring(entry))
164             f.close()
165
166             # Set last modified time for entry
167             os.utime(filename, (f_time, f_time))
168
169         time.sleep(5)
170
171
172 if __name__ == '__main__':
173     main()