Remove some headers Apt sets before returning full requests from the server.
[quix0rs-apt-p2p.git] / apt_p2p / apt_p2p.py
1
2 """The main program code.
3
4 @var DHT_PIECES: the maximum number of pieces to store with our contact info
5     in the DHT
6 @var TORRENT_PIECES: the maximum number of pieces to store as a separate entry
7     in the DHT
8 @var download_dir: the name of the directory to use for downloaded files
9 @var peer_dir: the name of the directory to use for peer downloads
10 """
11
12 from binascii import b2a_hex
13 from urlparse import urlunparse
14 from urllib import unquote
15 import os, re, sha
16
17 from twisted.internet import defer, reactor, protocol
18 from twisted.web2 import server, http, http_headers, static
19 from twisted.python import log, failure
20 from twisted.python.filepath import FilePath
21
22 from interfaces import IDHT, IDHTStats
23 from apt_p2p_conf import config
24 from PeerManager import PeerManager
25 from HTTPServer import TopLevel
26 from MirrorManager import MirrorManager
27 from CacheManager import CacheManager
28 from Hash import HashObject
29 from db import DB
30 from stats import StatsLogger
31 from util import findMyIPAddr, compact
32
33 DHT_PIECES = 4
34 TORRENT_PIECES = 70
35
36 download_dir = 'cache'
37 peer_dir = 'peers'
38
39 class AptP2P(protocol.Factory):
40     """The main code object that does all of the work.
41     
42     Contains all of the sub-components that do all the low-level work, and
43     coordinates communication between them.
44     
45     @type dhtClass: L{interfaces.IDHT}
46     @ivar dhtClass: the DHT class to use
47     @type cache_dir: L{twisted.python.filepath.FilePath}
48     @ivar cache_dir: the directory to use for storing all files
49     @type db: L{db.DB}
50     @ivar db: the database to use for tracking files and hashes
51     @type dht: L{interfaces.IDHT}
52     @ivar dht: the DHT instance
53     @type stats: L{stats.StatsLogger}
54     @ivar stats: the statistics logger to record sent data to
55     @type http_server: L{HTTPServer.TopLevel}
56     @ivar http_server: the web server that will handle all requests from apt
57         and from other peers
58     @type peers: L{PeerManager.PeerManager}
59     @ivar peers: the manager of all downloads from mirrors and other peers
60     @type mirrors: L{MirrorManager.MirrorManager}
61     @ivar mirrors: the manager of downloaded information about mirrors which
62         can be queried to get hashes from file names
63     @type cache: L{CacheManager.CacheManager}
64     @ivar cache: the manager of all downloaded files
65     @type my_contact: C{string}
66     @ivar my_contact: the 6-byte compact peer representation of this peer's
67         download information (IP address and port)
68     """
69     
70     def __init__(self, dhtClass):
71         """Initialize all the sub-components.
72         
73         @type dhtClass: L{interfaces.IDHT}
74         @param dhtClass: the DHT class to use
75         """
76         log.msg('Initializing the main apt_p2p application')
77         self.dhtClass = dhtClass
78
79     #{ Factory interface
80     def startFactory(self):
81         reactor.callLater(0, self._startFactory)
82         
83     def _startFactory(self):
84         log.msg('Starting the main apt_p2p application')
85         self.cache_dir = FilePath(config.get('DEFAULT', 'CACHE_DIR'))
86         if not self.cache_dir.child(download_dir).exists():
87             self.cache_dir.child(download_dir).makedirs()
88         if not self.cache_dir.child(peer_dir).exists():
89             self.cache_dir.child(peer_dir).makedirs()
90         self.db = DB(self.cache_dir.child('apt-p2p.db'))
91         self.dht = self.dhtClass()
92         self.dht.loadConfig(config, config.get('DEFAULT', 'DHT'))
93         self.dht.join().addCallbacks(self.joinComplete, self.joinError)
94         self.stats = StatsLogger(self.db)
95         self.http_server = TopLevel(self.cache_dir.child(download_dir), self.db, self)
96         self.http_server.getHTTPFactory().startFactory()
97         self.peers = PeerManager(self.cache_dir.child(peer_dir), self.dht, self.stats)
98         self.mirrors = MirrorManager(self.cache_dir)
99         self.cache = CacheManager(self.cache_dir.child(download_dir), self.db, self)
100         self.my_contact = None
101         
102     def stopFactory(self):
103         log.msg('Stoppping the main apt_p2p application')
104         self.http_server.getHTTPFactory().stopFactory()
105         self.stats.save()
106         self.db.close()
107     
108     def buildProtocol(self, addr):
109         return self.http_server.getHTTPFactory().buildProtocol(addr)
110         
111     #{ DHT Maintenance
112     def joinComplete(self, result):
113         """Complete the DHT join process and determine our download information.
114         
115         Called by the DHT when the join has been completed with information
116         on the external IP address and port of this peer.
117         """
118         my_addr = findMyIPAddr(result,
119                                config.getint(config.get('DEFAULT', 'DHT'), 'PORT'),
120                                config.getboolean('DEFAULT', 'LOCAL_OK'))
121         if not my_addr:
122             raise RuntimeError, "IP address for this machine could not be found"
123         self.my_contact = compact(my_addr, config.getint('DEFAULT', 'PORT'))
124         self.cache.scanDirectories()
125         reactor.callLater(60, self.refreshFiles)
126
127     def joinError(self, failure):
128         """Joining the DHT has failed."""
129         log.msg("joining DHT failed miserably")
130         log.err(failure)
131         raise RuntimeError, "IP address for this machine could not be found"
132     
133     def refreshFiles(self):
134         """Refresh any files in the DHT that are about to expire."""
135         expireAfter = config.gettime('DEFAULT', 'KEY_REFRESH')
136         hashes = self.db.expiredHashes(expireAfter)
137         if len(hashes.keys()) > 0:
138             log.msg('Refreshing the keys of %d DHT values' % len(hashes.keys()))
139         self._refreshFiles(None, hashes)
140         
141     def _refreshFiles(self, result, hashes):
142         if result is not None:
143             log.msg('Storage resulted in: %r' % result)
144
145         if hashes:
146             raw_hash = hashes.keys()[0]
147             self.db.refreshHash(raw_hash)
148             hash = HashObject(raw_hash, pieces = hashes[raw_hash]['pieces'])
149             del hashes[raw_hash]
150             storeDefer = self.store(hash)
151             storeDefer.addBoth(self._refreshFiles, hashes)
152         else:
153             reactor.callLater(60, self.refreshFiles)
154     
155     def getStats(self):
156         """Retrieve and format the statistics for the program.
157         
158         @rtype: C{string}
159         @return: the formatted HTML page containing the statistics
160         """
161         out = '<html><body>\n\n'
162         out += self.stats.formatHTML(self.my_contact)
163         out += '\n\n'
164         if IDHTStats.implementedBy(self.dhtClass):
165             out += self.dht.getStats()
166         out += '\n</body></html>\n'
167         return out
168
169     #{ Main workflow
170     def check_freshness(self, req, url, modtime, resp):
171         """Send a HEAD to the mirror to check if the response from the cache is still valid.
172         
173         @type req: L{twisted.web2.http.Request}
174         @param req: the initial request sent to the HTTP server by apt
175         @param url: the URI of the actual mirror request
176         @type modtime: C{int}
177         @param modtime: the modified time of the cached file (seconds since epoch)
178         @type resp: L{twisted.web2.http.Response}
179         @param resp: the response from the cache to be sent to apt
180         @rtype: L{twisted.internet.defer.Deferred}
181         @return: a deferred that will be called back with the correct response
182         """
183         log.msg('Checking if %s is still fresh' % url)
184         d = self.peers.get('', url, method = "HEAD", modtime = modtime)
185         d.addCallbacks(self.check_freshness_done, self.check_freshness_error,
186                        callbackArgs = (req, url, resp), errbackArgs = (req, url))
187         return d
188     
189     def check_freshness_done(self, resp, req, url, orig_resp):
190         """Process the returned response from the mirror.
191         
192         @type resp: L{twisted.web2.http.Response}
193         @param resp: the response from the mirror to the HEAD request
194         @type req: L{twisted.web2.http.Request}
195         @param req: the initial request sent to the HTTP server by apt
196         @param url: the URI of the actual mirror request
197         @type orig_resp: L{twisted.web2.http.Response}
198         @param orig_resp: the response from the cache to be sent to apt
199         """
200         if resp.code == 304:
201             log.msg('Still fresh, returning: %s' % url)
202             return orig_resp
203         else:
204             log.msg('Stale, need to redownload: %s' % url)
205             return self.get_resp(req, url)
206     
207     def check_freshness_error(self, err, req, url):
208         """Mirror request failed, continue with download.
209         
210         @param err: the response from the mirror to the HEAD request
211         @type req: L{twisted.web2.http.Request}
212         @param req: the initial request sent to the HTTP server by apt
213         @param url: the URI of the actual mirror request
214         """
215         log.err(err)
216         return self.get_resp(req, url)
217     
218     def get_resp(self, req, url):
219         """Lookup a hash for the file in the local mirror info.
220         
221         Starts the process of getting a response to an uncached apt request.
222         
223         @type req: L{twisted.web2.http.Request}
224         @param req: the initial request sent to the HTTP server by apt
225         @param url: the URI of the actual mirror request
226         @rtype: L{twisted.internet.defer.Deferred}
227         @return: a deferred that will be called back with the response
228         """
229         d = defer.Deferred()
230         
231         log.msg('Trying to find hash for %s' % url)
232         findDefer = self.mirrors.findHash(unquote(url))
233         
234         findDefer.addCallbacks(self.findHash_done, self.findHash_error, 
235                                callbackArgs=(req, url, d), errbackArgs=(req, url, d))
236         findDefer.addErrback(log.err)
237         return d
238     
239     def findHash_error(self, failure, req, url, d):
240         """Process the error in hash lookup by returning an empty L{HashObject}."""
241         log.err(failure)
242         self.findHash_done(HashObject(), req, url, d)
243         
244     def findHash_done(self, hash, req, url, d):
245         """Use the returned hash to lookup  the file in the cache.
246         
247         If the hash was not found, the workflow skips down to download from
248         the mirror (L{lookupHash_done}).
249         
250         @type hash: L{Hash.HashObject}
251         @param hash: the hash object containing the expected hash for the file
252         """
253         if hash.expected() is None:
254             log.msg('Hash for %s was not found' % url)
255             self.lookupHash_done([], req, hash, url, d)
256         else:
257             log.msg('Found hash %s for %s' % (hash.hexexpected(), url))
258             
259             # Lookup hash in cache
260             locations = self.db.lookupHash(hash.expected(), filesOnly = True)
261             self.getCachedFile(hash, req, url, d, locations)
262
263     def getCachedFile(self, hash, req, url, d, locations):
264         """Try to return the file from the cache, otherwise move on to a DHT lookup.
265         
266         @type locations: C{list} of C{dictionary}
267         @param locations: the files in the cache that match the hash,
268             the dictionary contains a key 'path' whose value is a
269             L{twisted.python.filepath.FilePath} object for the file.
270         """
271         if not locations:
272             log.msg('Failed to return file from cache: %s' % url)
273             self.lookupHash(req, hash, url, d)
274             return
275         
276         # Get the first possible location from the list
277         file = locations.pop(0)['path']
278         log.msg('Returning cached file: %s' % file.path)
279         
280         # Get it's response
281         resp = static.File(file.path).renderHTTP(req)
282         if isinstance(resp, defer.Deferred):
283             resp.addBoth(self._getCachedFile, hash, req, url, d, locations)
284         else:
285             self._getCachedFile(resp, hash, req, url, d, locations)
286         
287     def _getCachedFile(self, resp, hash, req, url, d, locations):
288         """Check the returned response to be sure it is valid."""
289         if isinstance(resp, failure.Failure):
290             log.msg('Got error trying to get cached file')
291             log.err()
292             # Try the next possible location
293             self.getCachedFile(hash, req, url, d, locations)
294             return
295             
296         log.msg('Cached response: %r' % resp)
297         
298         if resp.code >= 200 and resp.code < 400:
299             d.callback(resp)
300         else:
301             # Try the next possible location
302             self.getCachedFile(hash, req, url, d, locations)
303
304     def lookupHash(self, req, hash, url, d):
305         """Lookup the hash in the DHT."""
306         log.msg('Looking up hash in DHT for file: %s' % url)
307         key = hash.expected()
308         lookupDefer = self.dht.getValue(key)
309         lookupDefer.addBoth(self.lookupHash_done, req, hash, url, d)
310
311     def lookupHash_done(self, values, req, hash, url, d):
312         """Start the download of the file.
313         
314         The download will be from peers if the DHT lookup succeeded, or
315         from the mirror otherwise.
316         
317         @type values: C{list} of C{dictionary}
318         @param values: the returned values from the DHT containing peer
319             download information
320         """
321         # Remove some headers Apt sets in the request
322         req.headers.removeHeader('If-Modified-Since')
323         req.headers.removeHeader('Range')
324         req.headers.removeHeader('If-Range')
325         
326         if not isinstance(values, list) or not values:
327             if not isinstance(values, list):
328                 log.msg('DHT lookup for %s failed with error %r' % (url, values))
329             else:
330                 log.msg('Peers for %s were not found' % url)
331             getDefer = self.peers.get(hash, url)
332             getDefer.addCallback(self.cache.save_file, hash, url)
333             getDefer.addErrback(self.cache.save_error, url)
334             getDefer.addCallbacks(d.callback, d.errback)
335         else:
336             log.msg('Found peers for %s: %r' % (url, values))
337             # Download from the found peers
338             getDefer = self.peers.get(hash, url, values)
339             getDefer.addCallback(self.check_response, hash, url)
340             getDefer.addCallback(self.cache.save_file, hash, url)
341             getDefer.addErrback(self.cache.save_error, url)
342             getDefer.addCallbacks(d.callback, d.errback)
343             
344     def check_response(self, response, hash, url):
345         """Check the response from peers, and download from the mirror if it is not."""
346         if response.code < 200 or response.code >= 300:
347             log.msg('Download from peers failed, going to direct download: %s' % url)
348             getDefer = self.peers.get(hash, url)
349             return getDefer
350         return response
351         
352     def new_cached_file(self, file_path, hash, new_hash, url = None, forceDHT = False):
353         """Add a newly cached file to the mirror info and/or the DHT.
354         
355         If the file was downloaded, set url to the path it was downloaded for.
356         Doesn't add a file to the DHT unless a hash was found for it
357         (but does add it anyway if forceDHT is True).
358         
359         @type file_path: L{twisted.python.filepath.FilePath}
360         @param file_path: the location of the file in the local cache
361         @type hash: L{Hash.HashObject}
362         @param hash: the original (expected) hash object containing also the
363             hash of the downloaded file
364         @type new_hash: C{boolean}
365         @param new_hash: whether the has was new to this peer, and so should
366             be added to the DHT
367         @type url: C{string}
368         @param url: the URI of the location of the file in the mirror
369             (optional, defaults to not adding the file to the mirror info)
370         @type forceDHT: C{boolean}
371         @param forceDHT: whether to force addition of the file to the DHT
372             even if the hash was not found in a mirror
373             (optional, defaults to False)
374         """
375         if url:
376             self.mirrors.updatedFile(url, file_path)
377         
378         if self.my_contact and hash and new_hash and (hash.expected() is not None or forceDHT):
379             return self.store(hash)
380         return None
381             
382     def store(self, hash):
383         """Add a key/value pair for the file to the DHT.
384         
385         Sets the key and value from the hash information, and tries to add
386         it to the DHT.
387         """
388         key = hash.digest()
389         value = {'c': self.my_contact}
390         pieces = hash.pieceDigests()
391         
392         # Determine how to store any piece data
393         if len(pieces) <= 1:
394             pass
395         elif len(pieces) <= DHT_PIECES:
396             # Short enough to be stored with our peer contact info
397             value['t'] = {'t': ''.join(pieces)}
398         elif len(pieces) <= TORRENT_PIECES:
399             # Short enough to be stored in a separate key in the DHT
400             value['h'] = sha.new(''.join(pieces)).digest()
401         else:
402             # Too long, must be served up by our peer HTTP server
403             value['l'] = sha.new(''.join(pieces)).digest()
404
405         storeDefer = self.dht.storeValue(key, value)
406         storeDefer.addCallbacks(self.store_done, self.store_error,
407                                 callbackArgs = (hash, ), errbackArgs = (hash.digest(), ))
408         return storeDefer
409
410     def store_done(self, result, hash):
411         """Add a key/value pair for the pieces of the file to the DHT (if necessary)."""
412         log.msg('Added %s to the DHT: %r' % (hash.hexdigest(), result))
413         pieces = hash.pieceDigests()
414         if len(pieces) > DHT_PIECES and len(pieces) <= TORRENT_PIECES:
415             # Add the piece data key and value to the DHT
416             key = sha.new(''.join(pieces)).digest()
417             value = {'t': ''.join(pieces)}
418
419             storeDefer = self.dht.storeValue(key, value)
420             storeDefer.addCallbacks(self.store_torrent_done, self.store_error,
421                                     callbackArgs = (key, ), errbackArgs = (key, ))
422             return storeDefer
423         return result
424
425     def store_torrent_done(self, result, key):
426         """Adding the file to the DHT is complete, and so is the workflow."""
427         log.msg('Added torrent string %s to the DHT: %r' % (b2a_hex(key), result))
428         return result
429
430     def store_error(self, err, key):
431         """Adding to the DHT failed."""
432         log.msg('An error occurred adding %s to the DHT: %r' % (b2a_hex(key), err))
433         return err
434