forked from iambus/xunlei-lixian
-
Notifications
You must be signed in to change notification settings - Fork 0
/
lixian_tasks.py
322 lines (293 loc) · 9.01 KB
/
lixian_tasks.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
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
__all__ = ['search_tasks', 'find_task_by_url', 'find_task_by_url_or_path', 'find_tasks_to_download', 'find_torrent_tasks_to_download', 'find_normal_tasks_to_download', 'expand_bt_sub_tasks', 'is_url', 'is_local_bt']
import re
import os
import urllib2
import fileinput
from lixian_encoding import default_encoding
import lixian_hash_bt
import lixian_hash_ed2k
import lixian_extend_links
def native_to_utf_8(url):
try:
return url.decode(default_encoding).encode('utf-8')
except:
return url
def link_normalize(url):
from lixian_url import url_unmask, normalize_unicode_link
url = url_unmask(url)
if url.startswith('magnet:'):
return 'bt://'+lixian_hash_bt.magnet_to_infohash(url).encode('hex')
elif url.startswith('ed2k://'):
return lixian_hash_ed2k.parse_ed2k_id(url)
elif url.startswith('bt://'):
return url.lower()
elif url.startswith('http://') or url.startswith('ftp://'):
return normalize_unicode_link(url)
return url
def link_equals(x1, x2):
return link_normalize(x1) == link_normalize(x2)
def link_in(url, links):
for link in links:
if link_equals(url, link):
return True
def is_url(url):
return re.match(r'\w+://|magnet:', url)
def is_local_bt(url):
return (not is_url(url)) and url.lower().endswith('.torrent') and os.path.exists(url)
def is_id(x):
return re.match(r'^#?\d+(/[-.\w\[\],\s*]+)?$', x) or re.match(r'^#?\d+-\d+$', x)
def find_task_by_url(tasks, url):
for t in tasks:
if link_equals(t['original_url'], url):
return t
def find_task_by_url_or_path(tasks, url):
if is_url(url):
return find_task_by_url(tasks, url)
elif is_local_bt(url):
return find_task_by_url(tasks, 'bt://' + lixian_hash_bt.info_hash(url))
else:
raise NotImplementedError()
def find_tasks_by_range(tasks, x):
m = re.match(r'^#?(\d+)-(\d+)$', x)
begin = int(m.group(1))
end = int(m.group(2))
if begin <= end:
return filter(lambda x: begin <= x['#'] <= end, tasks)
else:
return reversed(filter(lambda x: end <= x['#'] <= begin, tasks))
def find_task_by_id(tasks, id):
for t in tasks:
if str(t['id']) == id or str(t['#']) == id or '#'+str(t['#']) == id:
return t
def find_tasks_by_id(tasks, id):
if re.match(r'^#?\d+-\d+$', id):
return find_tasks_by_range(tasks, id)
task_id, sub_id = re.match(r'^(#?\d+)(?:/([-.\w\[\],\s*]+))?$', id).groups()
task = find_task_by_id(tasks, task_id)
if not task:
return []
if not sub_id:
return [task]
assert task['type'] == 'bt', 'task %s is not a bt task' % task['name'].encode(default_encoding)
matched = []
if re.match(r'\[.*\]', sub_id):
for sub_id in re.split(r'\s*,\s*', sub_id[1:-1]):
assert re.match(r'^\d+(-\d+)?|\.\w+$', sub_id), sub_id
if sub_id.startswith('.'):
t = dict(task)
t['index'] = sub_id
matched.append(t)
elif '-' in sub_id:
start, end = map(int, sub_id.split('-'))
r = range(start, end+1) if start <= end else reversed(range(end, start+1))
for i in r:
t = dict(task)
t['index'] = str(i)
matched.append(t)
else:
assert re.match(r'^\d+$', sub_id), sub_id
t = dict(task)
t['index'] = sub_id
matched.append(t)
elif re.match(r'^\.\w+$', sub_id):
t = dict(task)
t['index'] = sub_id
matched.append(t)
elif sub_id == '*':
t = dict(task)
t['index'] = sub_id
matched.append(t)
else:
assert re.match(r'^\d+$', sub_id), sub_id
t = dict(task)
t['index'] = sub_id
matched.append(t)
return matched
def search_in_tasks(tasks, keywords):
found = []
for x in keywords:
# search url and local bt
if is_url(x) or is_local_bt(x):
task = find_task_by_url_or_path(tasks, x)
if task:
found.append(task)
else:
found.append(x) # keep the task order per arguments
continue
# search id
if is_id(x):
matched = find_tasks_by_id(tasks, x)
if matched:
found += matched
continue
# search date
if re.match(r'^\d{4}\.\d{2}\.\d{2}$', x):
raise NotImplementedError()
matched = filter(lambda t: t['date'] == v, tasks)
if matched:
found += matched
continue
# search name
if type(x) == str:
x = x.decode(default_encoding)
matched = filter(lambda t: t['name'].lower().find(x.lower()) != -1, tasks)
if matched:
found += matched
else:
# keyword not matched
pass
found = merge_bt_sub_tasks(found)
return filter(lambda x: type(x) == dict, found), filter(lambda x: type(x) != dict, found), found
def search_tasks(client, args, status='all'):
if status == 'all':
tasks = client.read_all_tasks()
elif status == 'completed':
tasks = client.read_all_tasks()
elif status == 'deleted':
tasks = client.read_all_deleted()
elif status == 'expired':
tasks = client.read_all_expired()
else:
raise NotImplementedError()
return search_in_tasks(tasks, list(args))[0]
def find_torrent_tasks_to_download(client, links):
tasks = client.read_all_tasks()
hashes = set(t['bt_hash'].lower() for t in tasks if t['type'] == 'bt')
link_hashes = []
for link in links:
if re.match(r'^(?:bt://)?([a-fA-F0-9]{40})$', link):
info_hash = link[-40:].lower()
if info_hash not in hashes:
print 'Adding bt task', link
client.add_torrent_task_by_info_hash(info_hash)
link_hashes.append(info_hash)
elif re.match(r'http://', link):
print 'Downloading torrent file from', link
torrent = urllib2.urlopen(link, timeout=60).read()
info_hash = lixian_hash_bt.info_hash_from_content(torrent)
if info_hash not in hashes:
print 'Adding bt task', link
client.add_torrent_task_by_content(torrent, os.path.basename(link))
link_hashes.append(info_hash)
elif os.path.exists(link):
with open(link, 'rb') as stream:
torrent = stream.read()
info_hash = lixian_hash_bt.info_hash_from_content(torrent)
if info_hash not in hashes:
print 'Adding bt task', link
client.add_torrent_task_by_content(torrent, os.path.basename(link))
link_hashes.append(info_hash)
else:
raise NotImplementedError('Unknown torrent '+link)
all_tasks = client.read_all_tasks()
tasks = []
for h in link_hashes:
for t in all_tasks:
if t['bt_hash'].lower() == h.lower():
tasks.append(t)
break
else:
raise NotImplementedError('not task found')
return tasks
def return_my_tasks(all_tasks, links):
tasks = []
for x in links:
if type(x) == dict:
tasks.append(x)
else:
task = find_task_by_url_or_path(all_tasks, x)
if not task:
raise NotImplementedError('task not found, wired: '+x)
tasks.append(task)
return tasks
def find_normal_tasks_to_download(client, links):
links = lixian_extend_links.extend_links(links)
all_tasks = client.read_all_tasks()
found, missing, all = search_in_tasks(all_tasks, links)
to_add = set(missing)
if to_add:
print 'Adding below tasks:'
for link in missing:
print link
links_to_add = filter(is_url, to_add)
if links_to_add:
client.add_batch_tasks(map(native_to_utf_8, links_to_add))
for link in to_add:
if is_url(link):
# add_batch_tasks doesn't work for bt task, add bt task one by one...
if link.startswith('bt://') or link.startswith('magnet:'):
client.add_task(link)
elif is_local_bt(link):
with open(link, 'rb') as stream:
torrent = stream.read()
client.add_torrent_task_by_content(torrent, os.path.basename(link))
else:
raise NotImplementedError('Unsupported: '+link)
all_tasks = client.read_all_tasks()
try:
return return_my_tasks(all_tasks, all)
except NotImplementedError:
import time
time.sleep(5)
return return_my_tasks(client.read_all_tasks(), all)
def find_tasks_to_download(client, args):
links = []
links.extend(args)
if args.input:
links.extend(line.strip() for line in fileinput.input(args.input) if line.strip())
if args.torrent:
return find_torrent_tasks_to_download(client, links)
else:
return find_normal_tasks_to_download(client, links)
def merge_bt_sub_tasks(tasks):
result_tasks = []
task_mapping = {}
for task in tasks:
if type(task) == dict:
id = task['id']
if id in task_mapping:
if 'index' in task and 'files' in task_mapping[id]:
task_mapping[id]['files'].append(task['index'])
else:
if 'index' in task:
t = dict(task)
t['files'] = [t['index']]
del t['index']
result_tasks.append(t)
task_mapping[id] = t
else:
result_tasks.append(task)
task_mapping[id] = task
else:
if task in task_mapping:
pass
else:
result_tasks.append(task)
task_mapping[task] = task
return result_tasks
def expand_bt_sub_tasks(client, task):
files = client.list_bt(task)
not_ready = []
single_file = False
if len(files) == 1 and files[0]['name'] == task['name']:
single_file = True
if 'files' in task:
ordered_files = []
indexed_files = dict((f['index'], f) for f in files)
subs = []
for index in task['files']:
if index == '*':
subs.extend([x['index'] for x in files])
elif index.startswith('.'):
subs.extend([x['index'] for x in files if x['name'].lower().endswith(index.lower())])
else:
subs.append(int(index))
for index in subs:
t = indexed_files[index]
if t not in ordered_files:
if t['status_text'] != 'completed':
not_ready.append(t)
else:
ordered_files.append(t)
files = ordered_files
return files, not_ready, single_file