· 5 years ago · Mar 01, 2021, 08:40 PM
1#!/usr/bin/python3
2
3import os
4import shutil
5import yaml
6import argparse
7import logging
8import logging.handlers
9from qbittorrentapi import Client
10import urllib3
11from collections import Counter
12
13# import apprise
14
15parser = argparse.ArgumentParser('qBittorrent Manager.',
16 description='A mix of scripts combined for managing qBittorrent.')
17parser.add_argument('-c', '--config-file',
18 dest='config',
19 action='store',
20 default='config.yml',
21 help='This is used if you want to use a different name for your config.yml. Example: tv.yml')
22parser.add_argument('-l', '--log-file',
23 dest='logfile',
24 action='store',
25 default='activity.log',
26 help='This is used if you want to use a different name for your log file. Example: tv.log')
27parser.add_argument('-m', '--manage',
28 dest='manage',
29 action='store_const',
30 const='manage',
31 help='Use this if you would like to update your tags AND'
32 ' categories AND remove unregistered torrents.')
33parser.add_argument('-s', '--cross-seed',
34 dest='cross_seed',
35 action='store_const',
36 const='cross_seed',
37 help='Use this after running cross-seed script to organize your torrents into specified '
38 'watch folders.')
39parser.add_argument('-g', '--cat-update',
40 dest='cat_update',
41 action='store_const',
42 const='cat_update',
43 help='Use this if you would like to update your categories.')
44parser.add_argument('-t', '--tag-update',
45 dest='tag_update',
46 action='store_const',
47 const='tag_update',
48 help='Use this if you would like to update your tags.')
49parser.add_argument('-r', '--rem-unregistered',
50 dest='rem_unregistered',
51 action='store_const',
52 const='rem_unregistered',
53 help='Use this if you would like to remove unregistered torrents.')
54parser.add_argument('--dry-run',
55 dest='dry_run',
56 action='store_const',
57 const='dry_run',
58 help='If you would like to see what is gonna happen but not actually delete or '
59 'tag/categorize anything.')
60parser.add_argument('--log',
61 dest='loglevel',
62 action='store',
63 default='INFO',
64 help='Change your log level. ')
65args = parser.parse_args()
66
67with open(args.config, 'r') as cfg_file:
68 cfg = yaml.load(cfg_file, Loader=yaml.FullLoader)
69
70urllib3.disable_warnings()
71
72file_name_format = args.logfile
73msg_format = '%(asctime)s - %(levelname)s: %(message)s'
74max_bytes = 1024 * 1024 * 2
75backup_count = 5
76
77logger = logging.getLogger('qBit Manage')
78logging.DRYRUN = 25
79logging.addLevelName(logging.DRYRUN, 'DRY-RUN')
80setattr(logger, 'dryrun', lambda dryrun, *args: logger._log(logging.DRYRUN, dryrun, args))
81log_lev = getattr(logging, args.loglevel.upper())
82logger.setLevel(log_lev)
83
84file_handler = logging.handlers.RotatingFileHandler(filename=file_name_format,
85 maxBytes=max_bytes,
86 backupCount=backup_count)
87file_handler.setLevel(log_lev)
88file_formatter = logging.Formatter(msg_format)
89file_handler.setFormatter(file_formatter)
90logger.addHandler(file_handler)
91
92stream_handler = logging.StreamHandler()
93stream_handler.setLevel(log_lev)
94stream_formatter = logging.Formatter(msg_format)
95stream_handler.setFormatter(stream_formatter)
96logger.addHandler(stream_handler)
97
98# Actual API call to connect to qbt.
99host = cfg['qbt']['host']
100if 'user' in cfg['qbt']:
101 username = cfg['qbt']['user']
102else:
103 username = ''
104if 'pass' in cfg['qbt']:
105 password = cfg['qbt']['pass']
106else:
107 password = ''
108
109client = Client(host=host,
110 username=username,
111 password=password)
112
113
114def convert_bytes(bytes_number):
115 tags = ["Byte", "Kilobyte", "Megabyte", "Gigabyte", "Terabyte"]
116 i = 0
117 double_bytes = bytes_number
118 while i < len(tags) and bytes_number >= 1024:
119 double_bytes = bytes_number / 1024.0
120 i = i + 1
121 bytes_number = bytes_number / 1024
122 return str(round(double_bytes, 2)) + " " + tags[i]
123
124
125def trunc_val(s, d, n=3):
126 return d.join(s.split(d, n)[:n])
127
128
129def get_category(path):
130 for cat, attr in cfg['cat'].items():
131 for attr_path in attr:
132 if 'save_path' in attr_path and attr_path['save_path'] in path:
133 category = cat
134 return category
135 else:
136 category = ''
137 logger.warning('No categories matched. Check your config.yml file. - Setting category to NULL')
138 return category
139
140
141def get_tags(url):
142 tag_path = cfg['tags']
143 for i, f in tag_path.items():
144 if i in url:
145 tag = f
146 return tag
147 else:
148 tag = ''
149 logger.warning('No tags matched. Check your config.yml file. Setting tag to NULL')
150 return tag
151
152
153def get_name(t_list):
154 dupes = []
155 no_dupes = []
156 t_name = [torrent.name for torrent in t_list]
157 dupes = [s for s in t_name if t_name.count(s) > 1 if s not in dupes]
158 no_dupes = [s for s in t_name if t_name.count(s) == 1 if s not in no_dupes]
159 return dupes, no_dupes
160
161
162# Will create a 2D Dictionary with the torrent name as the key
163# torrentdict = {'TorrentName1' : {'Category':'TV', 'save_path':'/data/torrents/TV'},
164# 'TorrentName2' : {'Category':'Movies', 'save_path':'/data/torrents/Movies'}}
165def get_torrent_info(t_list):
166 torrentdict = {}
167 for torrent in t_list:
168 save_path = torrent.save_path
169 torrent_hash = torrent.hash # Needed to pause and modify state of torrent.
170 amount_left = torrent.amount_left # This is 0 if complete. Can use this or torrent.progress depends on how spunky you're feeling :P
171 category = get_category(save_path)
172 total_size = torrent.total_size # Number in bytes of size of torrent can use this to edit priority based on size
173 torrent_state = torrent.state # Gives state of torrent but can use these instead https://qbittorrent-api.readthedocs.io/en/latest/apidoc/torrent_states.html
174 torrent_progress = torrent.progress # This is in a float format Example: 1 = 100%, 0.74 = 74% etc.
175 priority = torrent.priority # Gives priority number of torrent.
176 torrentattr = {'Category': category,
177 'save_path': save_path,
178 'torrent_hash': torrent_hash,
179 'amount_left': amount_left,
180 'torrent_state': torrent_state,
181 'torrent_progress': torrent_progress,
182 'total_size': total_size,
183 'priority': priority}
184 torrentdict[torrent.name] = torrentattr
185 return torrentdict
186
187
188def get_hash_size(t_list):
189 t_hash = []
190 t_size = []
191 for torrent in t_list:
192 if torrent.state_enum.is_checking:
193 t_hash.append(torrent.hash)
194 t_size.append(torrent.total_size)
195 t_size.sort()
196 return t_hash, t_size
197
198
199# Function used to move any torrents from the cross seed directory to the correct save directory
200def cross_seed():
201 if args.cross_seed == 'cross_seed':
202 categories = []
203 total = 0
204 torrents_moved = ''
205 cs_files = [f for f in os.listdir(os.path.join(cfg['directory']['cross_seed'], '')) if f.endswith('torrent')]
206 dir_cs = os.path.join(cfg['directory']['cross_seed'], '')
207 torrent_list = client.torrents.info()
208 torrentdict = get_torrent_info(torrent_list)
209 for file in cs_files:
210 t_name = file.split(']', 2)[2].split('.torrent')[0]
211 torrentdict_file = dict(filter(lambda item: t_name in item[0], torrentdict.items()))
212 if torrentdict_file:
213 t_name = next(iter(torrentdict_file))
214 category = torrentdict[t_name]['Category']
215 dest = os.path.join(torrentdict[t_name]['save_path'], '')
216 for attr_path in cfg['cat'][category]:
217 if 'watch_path' in attr_path:
218 dest = os.path.join(attr_path['watch_path'], '')
219 if 'remote_dir' in cfg:
220 for dir in cfg['remote_dir']:
221 if dir in dest:
222 dest = dest.replace(dir, cfg['remote_dir'][dir])
223 src = dir_cs + file
224 dest += file
225 categories.append(category)
226 s_path = torrentdict[t_name]['save_path']
227 if args.dry_run == 'dry_run':
228 logger.dryrun(f'Adding {t_name} to qBittorrent with: '
229 f'\n - Category: {category}'
230 f'\n - Save_Path: {s_path}'
231 f'\n - Paused: True')
232 else:
233 client.torrents.add(torrent_files=src,
234 save_path=torrentdict[t_name]['save_path'],
235 category=category,
236 is_paused=True)
237 logger.info(f'Adding {t_name} to qBittorrent with: '
238 f'\n - Category: {category}'
239 f'\n - Save_Path: {s_path}'
240 f'\n - Paused: True')
241 client.torrents.recheck(torrentdict[t_name]['torrent_hash'])
242 # shutil.move(src, dest)
243 logger.info(f'Moving {src} to {dest}')
244 else:
245 if args.dry_run == 'dry_run':
246 logger.dryrun(f'{t_name} not found in torrents.')
247 else:
248 logger.warning(f'{t_name} not found in torrents.')
249
250
251 for torrent in torrent_list:
252 if torrentdict[torrent.hash] == torrent.hash:
253
254
255
256 t_hash, t_size = get_hash_size(torrent_list)
257 for torrent in torrent_list:
258 for idx, i in enumerate(t_size, start=1):
259 if torrent.hash in t_hash:
260 if i == torrent.total_size:
261 i = convert_bytes(i)
262 print(f'\n Priority: {idx} '
263 f'\n File_size: {i} '
264 f'\n Hash: {torrent.hash}')
265 # torrent.file_priority(torrent_hash=torrent.hash, priority=idx)
266 numcategory = Counter(categories)
267 if args.dry_run == 'dry_run':
268 for c in numcategory:
269 total += numcategory[c]
270 torrents_moved += f'\n - {c} .torrents not moved: {numcategory[c]}'
271 torrents_moved += f'\n -- Total .torrents not moved: {total}'
272 logger.dryrun(torrents_moved)
273 else:
274 for c in numcategory:
275 total += numcategory[c]
276 torrents_moved += f'\n - {c} .torrents moved: {numcategory[c]}'
277 torrents_moved += f'\n -- Total .torrents moved: {total}'
278 logger.info(torrents_moved)
279
280
281def update_category():
282 if args.manage == 'manage' or args.cat_update == 'cat_update':
283 num_cat = 0
284 torrent_list = client.torrents.info()
285 for torrent in torrent_list:
286 if torrent.category == '':
287 for x in torrent.trackers:
288 if x.url.startswith('http'):
289 t_url = trunc_val(x.url, '/')
290 new_cat = get_category(torrent.save_path)
291 if args.dry_run == 'dry_run':
292 logger.dryrun(f'\n - Torrent Name: {torrent.name}'
293 f'\n - New Category: {new_cat}'
294 f'\n - Tracker: {t_url}')
295 num_cat += 1
296 else:
297 logger.info(f'\n - Torrent Name: {torrent.name}'
298 f'\n - New Category: {new_cat}'
299 f'\n - Tracker: {t_url}')
300 torrent.set_category(category=new_cat)
301 num_cat += 1
302 if args.dry_run == 'dry_run':
303 if num_cat >= 1:
304 logger.dryrun(f'Did not update {num_cat} new categories.')
305 else:
306 logger.dryrun(f'No new torrents to categorize.')
307 else:
308 if num_cat >= 1:
309 logger.info(f'Updated {num_cat} new categories.')
310 else:
311 logger.info(f'No new torrents to categorize.')
312
313
314def update_tags():
315 if args.manage == 'manage' or args.tag_update == 'tag_update':
316 num_tags = 0
317 torrent_list = client.torrents.info()
318 for torrent in torrent_list:
319 if torrent.tags == '':
320 for x in torrent.trackers:
321 if x.url.startswith('http'):
322 t_url = trunc_val(x.url, '/')
323 new_tag = get_tags(x.url)
324 if args.dry_run == 'dry_run':
325 logger.dryrun(f'\n - Torrent Name: {torrent.name}'
326 f'\n - New Tag: {new_tag}'
327 f'\n - Tracker: {t_url}')
328 num_tags += 1
329 else:
330 logger.info(f'\n - Torrent Name: {torrent.name}'
331 f'\n - New Tag: {new_tag}'
332 f'\n - Tracker: {t_url}')
333 torrent.add_tags(tags=new_tag)
334 num_tags += 1
335 if args.dry_run == 'dry_run':
336 if num_tags >= 1:
337 logger.dryrun(f'Did not update {num_tags} new tags.')
338 else:
339 logger.dryrun('No new torrents to tag.')
340 else:
341 if num_tags >= 1:
342 logger.info(f'Updated {num_tags} new tags.')
343 else:
344 logger.info('No new torrents to tag. ')
345
346
347def rem_unregistered():
348 if args.manage == 'manage' or args.rem_unregistered == 'rem_unregistered':
349 torrent_list = client.torrents.info()
350 dupes, no_dupes = get_name(torrent_list)
351 rem_unr = 0
352 del_tor = 0
353 for torrent in torrent_list:
354 for status in torrent.trackers:
355 for x in torrent.trackers:
356 if x.url.startswith('http'):
357 t_url = trunc_val(x.url, '/')
358 n_info = (f'\n - Torrent Name: {torrent.name} '
359 f'\n - Status: {status.msg} '
360 f'\n - Tracker: {t_url} '
361 f'\n - Deleted .torrent but not content files.')
362 n_d_info = (f'\n - Torrent Name: {torrent.name} '
363 f'\n - Status: {status.msg} '
364 f'\n - Tracker: {t_url} '
365 f'\n - Deleted .torrent AND content files.')
366 if 'Unregistered torrent' in status.msg or 'Torrent is not found' in status.msg:
367 if torrent.name in dupes:
368 if args.dry_run == 'dry_run':
369 logger.dryrun(n_info)
370 rem_unr += 1
371 else:
372 logger.info(n_info)
373 torrent.delete(hash=torrent.hash, delete_files=False)
374 rem_unr += 1
375 elif torrent.name in no_dupes:
376 if args.dry_run == 'dry_run':
377 logger.dryrun(n_d_info)
378 del_tor += 1
379 else:
380 logger.info(n_d_info)
381 torrent.delete(hash=torrent.hash, delete_files=True)
382 del_tor += 1
383 if args.dry_run == 'dry_run':
384 if rem_unr >= 1 or del_tor >= 1:
385 logger.dryrun(f'Did not delete {rem_unr} .torrents(s) or content files.')
386 logger.dryrun(f'Did not delete {del_tor} .torrents(s) or content files.')
387 else:
388 logger.dryrun('No unregistered torrents found.')
389 else:
390 if rem_unr >= 1 or del_tor >= 1:
391 logger.info(f'Deleted {rem_unr} .torrents(s) but not content files.')
392 logger.info(f'Deleted {del_tor} .torrents(s) AND content files.')
393 else:
394 logger.info('No unregistered torrents found.')
395
396
397def run():
398 update_category()
399 update_tags()
400 rem_unregistered()
401 cross_seed()
402
403
404if __name__ == '__main__':
405 run()