Deal with the case where socket.recv fails, e.g. connection reset by peer.
[cheatproxy.git] / cheatbt.py
blob81f918837433b4dbb07fbb18893ec8b57a114dbc
1 import logging
2 import urlparse
4 from odict import OrderedDict
6 def load_mappings(fname="trackers"):
7 """
8 Initialises tracker to ratio multiple mappings from the specified
9 file.
10 """
12 logger = logging.getLogger("cheatbt")
13 logger.info('loading mappings from "' + fname + '"')
15 f = open(fname)
16 mappings = {}
18 for l in f:
19 tracker, multiple = l.split(":")
20 mappings[tracker.strip()] = int(multiple.strip())
22 f.close()
24 if "default" not in mappings:
25 mappings["default"] = 1
27 return mappings
29 def cheat(url, mappings):
30 """
31 Modifies BitTorrent tracker URLs, faking the amount of data
32 uploaded. All other URLs should pass through unimpeded.
33 """
35 parsed = urlparse.urlparse(url)
37 if "=" not in parsed.query:
38 return url
40 query = OrderedDict([i.split("=") for i in parsed.query.split("&")])
41 if "uploaded" not in query or query["uploaded"] == "0":
42 return url
44 if parsed.hostname in mappings:
45 multiple = mappings[parsed.hostname]
46 else:
47 if "default" in mappings:
48 multiple = mappings["default"]
49 else:
50 return url
52 # Don't bother munging the URL if the upload amount isn't going
53 # to change.
54 if multiple == 1:
55 return url
57 fakeupload = int(query["uploaded"])
59 logger = logging.getLogger("cheatbt")
60 logger.debug("%s: %d -> %d" % (parsed.hostname, fakeupload,
61 fakeupload * multiple))
63 fakeupload = fakeupload * multiple
64 query["uploaded"] = str(fakeupload)
66 new_query = ""
67 for k in query.keys():
68 new_query += k + "=" + query[k] + "&"
69 new_query = new_query[:-1] # Remove trailing "&"
71 # <scheme>://<netloc>/<path>;<params>?<query>#<fragment>
72 new_url = urlparse.urlunparse((parsed.scheme, parsed.netloc, parsed.path,
73 parsed.params, new_query, parsed.fragment))
75 return new_url