-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
142 lines (108 loc) · 5.08 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
from struct import pack, unpack
from binascii import unhexlify
from bencode import bencode, bdecode, decode_dict
from time import sleep
from hashlib import sha1
from twisted.internet import protocol, defer
from twisted.protocols import policies
class BitTorrentClient(protocol.Protocol, policies.TimeoutMixin):
def __init__(self, info_hash, peer_id, on_metadata_loaded):
self._info_hash = info_hash
self._peer_id = peer_id
self._buffer = buffer("")
self._read_handshake = True
self._metadata = {}
self._deferred = defer.Deferred()
self._deferred.addCallback(on_metadata_loaded, info_hash)
@staticmethod
def parseMessage(message):
# Return message code and message data
if message:
return (unpack("B", message[:1])[0], message[1:])
else:
return None
def sendExtendedMessage(self, message_id, message_data):
buf = pack("BB", 20, message_id) + bencode(message_data)
self.transport.write(pack("!I", len(buf)) + buf)
def handleMessage(self, msg_code, msg_data):
if msg_code == 20:
# If he send extended message, we can extend connection time
self.resetTimeout()
# Extended handshake
if ord(msg_data[0]) == 0:
hs_data = bdecode(msg_data[1:])
if "metadata_size" in hs_data and "m" in hs_data and "ut_metadata" in hs_data["m"]:
metadata_size = hs_data["metadata_size"]
ut_metadata_id = hs_data["m"]["ut_metadata"]
hs_response = {"e": 0,
"metadata_size": hs_data["metadata_size"],
"v": "\xce\xbcTorrent 3.4.9",
"m": {"ut_metadata": 1},
"reqq": 255}
# Response extended handshake
self.sendExtendedMessage(0, hs_response)
sleep(0.5)
# Request metadata
for i in range(0, 1 + metadata_size / (16 * 1024)):
self.sendExtendedMessage(ut_metadata_id, {"msg_type": 0, "piece": i})
sleep(0.05)
else:
self.transport.loseConnection()
elif ord(msg_data[0]) == 1:
r, l = decode_dict(msg_data[1:], 0)
if r["msg_type"] == 1:
self._metadata[r["piece"]] = msg_data[l + 1:]
metadata = reduce(lambda r, e: r + self._metadata[e], sorted(self._metadata.keys()), "")
if len(metadata) == r["total_size"]:
if sha1(metadata).digest() == self._info_hash:
self._deferred.callback(bdecode(metadata))
# Abort connection anyway
self.transport.abortConnection()
def connectionMade(self):
# Set connection timeout in 10 seconds (after 10 seconds idle connection will be aborted)
self.setTimeout(10)
# Send handshake
bp = list("BitTorrent protocol")
self.transport.write(pack("B19c", 19, *bp))
self.transport.write(unhexlify("0000000000100005"))
self.transport.write(self._info_hash)
self.transport.write(self._peer_id)
def dataReceived(self, data):
self._buffer = buffer(self._buffer) + buffer(data)
if self._read_handshake:
if len(self._buffer) >= 68:
# Skip handshake response
self._buffer = self._buffer[68:]
self._read_handshake = False
else:
return
else:
# Read regular message
while self._buffer:
msg_len = unpack("!I", self._buffer[:4])[0]
if len(self._buffer) >= msg_len + 4:
message = self.parseMessage(self._buffer[4: msg_len + 4])
if message:
self.handleMessage(*message)
self._buffer = self._buffer[msg_len + 4:]
else:
break
def timeoutConnection(self):
self.transport.abortConnection()
class BitTorrentFactory(protocol.ClientFactory):
protocol = BitTorrentClient
def __init__(self, info_hash, peer_id, on_metadata_loaded):
self._info_hash = info_hash
self._peer_id = peer_id
self._on_metadata_loaded = on_metadata_loaded
def buildProtocol(self, addr):
p = self.protocol(self._info_hash, self._peer_id, self._on_metadata_loaded)
p.factory = self
return p
def print_metadata(metadata):
print metadata
factory = BitTorrentFactory(info_hash=unhexlify('E4DFB9BC728B5554F81CBF97637F7EA5151BF565'),
peer_id=unhexlify('cd2e6673b9f2a21cad1e605fe5fb745b9f7a214d'),
on_metadata_loaded=print_metadata)
reactor.connectTCP("127.0.0.1", 16762, factory)
reactor.run()