qbit_manage/qbit_manage.py
2021-03-06 22:52:19 -05:00

397 lines
17 KiB
Python

#!/usr/bin/python3
import os
import shutil
import yaml
import argparse
import logging
import logging.handlers
from qbittorrentapi import Client
import urllib3
from collections import Counter
# import apprise
parser = argparse.ArgumentParser('qBittorrent Manager.',
description='A mix of scripts combined for managing qBittorrent.')
parser.add_argument('-c', '--config-file',
dest='config',
action='store',
default='config.yml',
help='This is used if you want to use a different name for your config.yml. Example: tv.yml')
parser.add_argument('-l', '--log-file',
dest='logfile',
action='store',
default='activity.log',
help='This is used if you want to use a different name for your log file. Example: tv.log')
parser.add_argument('-m', '--manage',
dest='manage',
action='store_const',
const='manage',
help='Use this if you would like to update your tags AND'
' categories AND remove unregistered torrents.')
parser.add_argument('-s', '--cross-seed',
dest='cross_seed',
action='store_const',
const='cross_seed',
help='Use this after running cross-seed script to add your torrents into qBittorrent')
parser.add_argument('-re', '--recheck',
dest='recheck',
action='store_const',
const='recheck',
help='Recheck paused torrents sorted by lowest size. Resume if Completed.')
parser.add_argument('-g', '--cat-update',
dest='cat_update',
action='store_const',
const='cat_update',
help='Use this if you would like to update your categories.')
parser.add_argument('-t', '--tag-update',
dest='tag_update',
action='store_const',
const='tag_update',
help='Use this if you would like to update your tags.')
parser.add_argument('-r', '--rem-unregistered',
dest='rem_unregistered',
action='store_const',
const='rem_unregistered',
help='Use this if you would like to remove unregistered torrents.')
parser.add_argument('--dry-run',
dest='dry_run',
action='store_const',
const='dry_run',
help='If you would like to see what is gonna happen but not actually delete or '
'tag/categorize anything.')
parser.add_argument('--log',
dest='loglevel',
action='store',
default='INFO',
help='Change your log level. ')
args = parser.parse_args()
with open(args.config, 'r') as cfg_file:
cfg = yaml.load(cfg_file, Loader=yaml.FullLoader)
urllib3.disable_warnings()
file_name_format = args.logfile
msg_format = '%(asctime)s - %(levelname)s: %(message)s'
max_bytes = 1024 * 1024 * 2
backup_count = 5
logger = logging.getLogger('qBit Manage')
logging.DRYRUN = 25
logging.addLevelName(logging.DRYRUN, 'DRY-RUN')
setattr(logger, 'dryrun', lambda dryrun, *args: logger._log(logging.DRYRUN, dryrun, args))
log_lev = getattr(logging, args.loglevel.upper())
logger.setLevel(log_lev)
file_handler = logging.handlers.RotatingFileHandler(filename=file_name_format,
maxBytes=max_bytes,
backupCount=backup_count)
file_handler.setLevel(log_lev)
file_formatter = logging.Formatter(msg_format)
file_handler.setFormatter(file_formatter)
logger.addHandler(file_handler)
stream_handler = logging.StreamHandler()
stream_handler.setLevel(log_lev)
stream_formatter = logging.Formatter(msg_format)
stream_handler.setFormatter(stream_formatter)
logger.addHandler(stream_handler)
# Actual API call to connect to qbt.
host = cfg['qbt']['host']
if 'user' in cfg['qbt']:
username = cfg['qbt']['user']
else:
username = ''
if 'pass' in cfg['qbt']:
password = cfg['qbt']['pass']
else:
password = ''
client = Client(host=host,
username=username,
password=password)
def trunc_val(s, d, n=3):
return d.join(s.split(d, n)[:n])
def get_category(path):
cat_path = cfg["cat"]
for i, f in cat_path.items():
if f in path:
category = i
return category
else:
category = ''
logger.warning('No categories matched. Check your config.yml file. - Setting tag to NULL')
return category
def get_tags(url):
tag_path = cfg['tags']
for i, f in tag_path.items():
if i in url:
tag = f
return tag
else:
tag = ''
logger.warning('No tags matched. Check your config.yml file. Setting tag to NULL')
return tag
# Will create a 2D Dictionary with the torrent name as the key
# torrentdict = {'TorrentName1' : {'Category':'TV', 'save_path':'/data/torrents/TV', 'count':1, 'msg':'[]'},
# 'TorrentName2' : {'Category':'Movies', 'save_path':'/data/torrents/Movies'}, 'count':2, 'msg':'[]'}
def get_torrent_info(t_list):
torrentdict = {}
for torrent in t_list:
save_path = torrent.save_path
category = get_category(save_path)
is_complete = False
if torrent.name in torrentdict:
t_count = torrentdict[torrent.name]['count'] + 1
msg_list = torrentdict[torrent.name]['msg']
is_complete = True if torrentdict[torrent.name]['is_complete'] == True else torrent.state_enum.is_complete
else:
t_count = 1
msg_list = []
is_complete = torrent.state_enum.is_complete
msg = [x.msg for x in torrent.trackers if x.url.startswith('http')][0]
msg_list.append(msg)
torrentattr = {'Category': category, 'save_path': save_path, 'count': t_count, 'msg': msg_list, 'is_complete': is_complete}
torrentdict[torrent.name] = torrentattr
return torrentdict
# Function used to recheck paused torrents sorted by size and resume torrents that are completed
def recheck(torrentdict=None):
if args.cross_seed == 'cross_seed' or args.manage == 'manage' or args.recheck == 'recheck':
#sort by size and paused
torrent_sorted_list = client.torrents.info(status_filter='paused',sort='size')
if torrentdict is None: torrentdict = get_torrent_info(client.torrents.info(sort='added_on',reverse=True))
for torrent in torrent_sorted_list:
new_tag = [get_tags(x.url) for x in torrent.trackers if x.url.startswith('http')]
if torrent.tags == '': torrent.add_tags(tags=new_tag)
#Resume torrent if completed
if torrent.progress == 1:
if args.dry_run == 'dry_run':
logger.dryrun(f'\n - Not Resuming {new_tag} - {torrent.name}')
else:
logger.info(f'\n - Resuming {new_tag} - {torrent.name}')
torrent.resume()
#Recheck
elif torrent.progress == 0 and torrentdict[torrent.name]['is_complete']:
if args.dry_run == 'dry_run':
logger.dryrun(f'\n - Not Rechecking {new_tag} - {torrent.name}')
else:
logger.info(f'\n - Rechecking {new_tag} - {torrent.name}')
torrent.recheck()
# Function used to move any torrents from the cross seed directory to the correct save directory
def cross_seed():
if args.cross_seed == 'cross_seed':
# List of categories for all torrents moved
categories = []
# Keep track of total torrents moved
total = 0
# Used to output the final list torrents moved to output in the log
torrents_added = ''
# Only get torrent files
cs_files = [f for f in os.listdir(os.path.join(cfg['directory']['cross_seed'], '')) if f.endswith('torrent')]
dir_cs = os.path.join(cfg['directory']['cross_seed'], '')
dir_cs_out = os.path.join(dir_cs,'qbit_manage_added')
os.makedirs(dir_cs_out,exist_ok=True)
torrent_list = client.torrents.info(sort='added_on',reverse=True)
torrentdict = get_torrent_info(torrent_list)
for file in cs_files:
t_name = file.split(']', 2)[2].split('.torrent')[0]
# Substring Key match in dictionary (used because t_name might not match exactly with torrentdict key)
# Returned the dictionary of filtered item
torrentdict_file = dict(filter(lambda item: t_name in item[0], torrentdict.items()))
if torrentdict_file:
# Get the exact torrent match name from torrentdict
t_name = next(iter(torrentdict_file))
category = torrentdict[t_name]['Category']
dest = os.path.join(torrentdict[t_name]['save_path'], '')
src = os.path.join(dir_cs,file)
dir_cs_out = os.path.join(dir_cs,'qbit_manage_added',file)
categories.append(category)
if args.dry_run == 'dry_run':
logger.dryrun(f'Not Adding {t_name} to qBittorrent with: '
f'\n - Category: {category}'
f'\n - Save_Path: {dest}'
f'\n - Paused: True')
else:
if torrentdict[t_name]['is_complete']:
client.torrents.add(torrent_files=src,
save_path=dest,
category=category,
is_paused=True)
shutil.move(src, dir_cs_out)
logger.info(f'Adding {t_name} to qBittorrent with: '
f'\n - Category: {category}'
f'\n - Save_Path: {dest}'
f'\n - Paused: True')
else:
logger.info(f'Found {t_name} in {dir_cs} but original torrent is not complete. Not adding to qBittorrent')
else:
if args.dry_run == 'dry_run':
logger.dryrun(f'{t_name} not found in torrents.')
else:
logger.warning(f'{t_name} not found in torrents.')
recheck(torrentdict)
numcategory = Counter(categories)
if args.dry_run == 'dry_run':
for c in numcategory:
total += numcategory[c]
torrents_added += f'\n - {c} .torrents not added: {numcategory[c]}'
torrents_added += f'\n -- Total .torrents not added: {total}'
logger.dryrun(torrents_added)
else:
for c in numcategory:
total += numcategory[c]
torrents_added += f'\n - {c} .torrents added: {numcategory[c]}'
torrents_added += f'\n -- Total .torrents added: {total}'
logger.info(torrents_added)
def update_category():
if args.manage == 'manage' or args.cat_update == 'cat_update':
num_cat = 0
torrent_list = client.torrents.info(sort='added_on',reverse=True)
for torrent in torrent_list:
if torrent.category == '':
for x in torrent.trackers:
if x.url.startswith('http'):
t_url = trunc_val(x.url, '/')
new_cat = get_category(torrent.save_path)
if args.dry_run == 'dry_run':
logger.dryrun(f'\n - Torrent Name: {torrent.name}'
f'\n - New Category: {new_cat}'
f'\n - Tracker: {t_url}')
num_cat += 1
else:
logger.info(f'\n - Torrent Name: {torrent.name}'
f'\n - New Category: {new_cat}'
f'\n - Tracker: {t_url}')
torrent.set_category(category=new_cat)
num_cat += 1
if args.dry_run == 'dry_run':
if num_cat >= 1:
logger.dryrun(f'Did not update {num_cat} new categories.')
else:
logger.dryrun(f'No new torrents to categorize.')
else:
if num_cat >= 1:
logger.info(f'Updated {num_cat} new categories.')
else:
logger.info(f'No new torrents to categorize.')
def update_tags():
if args.manage == 'manage' or args.tag_update == 'tag_update':
num_tags = 0
torrent_list = client.torrents.info(sort='added_on',reverse=True)
for torrent in torrent_list:
if torrent.tags == '':
for x in torrent.trackers:
if x.url.startswith('http'):
t_url = trunc_val(x.url, '/')
new_tag = get_tags(x.url)
if args.dry_run == 'dry_run':
logger.dryrun(f'\n - Torrent Name: {torrent.name}'
f'\n - New Tag: {new_tag}'
f'\n - Tracker: {t_url}')
num_tags += 1
else:
logger.info(f'\n - Torrent Name: {torrent.name}'
f'\n - New Tag: {new_tag}'
f'\n - Tracker: {t_url}')
torrent.add_tags(tags=new_tag)
num_tags += 1
if args.dry_run == 'dry_run':
if num_tags >= 1:
logger.dryrun(f'Did not update {num_tags} new tags.')
else:
logger.dryrun('No new torrents to tag.')
else:
if num_tags >= 1:
logger.info(f'Updated {num_tags} new tags.')
else:
logger.info('No new torrents to tag. ')
def rem_unregistered():
if args.manage == 'manage' or args.rem_unregistered == 'rem_unregistered':
torrent_list = client.torrents.info(sort='added_on',reverse=True)
torrentdict = get_torrent_info(torrent_list)
rem_unr = 0
del_tor = 0
for torrent in torrent_list:
t_name = torrent.name
t_count = torrentdict[t_name]['count']
t_msg = torrentdict[t_name]['msg']
for x in torrent.trackers:
if x.url.startswith('http'):
t_url = trunc_val(x.url, '/')
n_info = (f'\n - Torrent Name: {t_name} '
f'\n - Status: {x.msg} '
f'\n - Tracker: {t_url} '
f'\n - Deleted .torrent but not content files.')
n_d_info = (f'\n - Torrent Name: {t_name} '
f'\n - Status: {x.msg} '
f'\n - Tracker: {t_url} '
f'\n - Deleted .torrent AND content files.')
if 'Unregistered torrent' in x.msg or 'Torrent is not found' in x.msg or 'Torrent not registered' in x.msg:
if t_count > 1:
if args.dry_run == 'dry_run':
if '' in t_msg:
logger.dryrun(n_info)
rem_unr += 1
else:
logger.dryrun(n_d_info)
del_tor += 1
else:
# Checks if any of the original torrents are working
if '' in t_msg:
logger.info(n_info)
torrent.delete(hash=torrent.hash, delete_files=False)
rem_unr += 1
else:
logger.info(n_d_info)
torrent.delete(hash=torrent.hash, delete_files=True)
del_tor += 1
else:
if args.dry_run == 'dry_run':
logger.dryrun(n_d_info)
del_tor += 1
else:
logger.info(n_d_info)
torrent.delete(hash=torrent.hash, delete_files=True)
del_tor += 1
if args.dry_run == 'dry_run':
if rem_unr >= 1 or del_tor >= 1:
logger.dryrun(f'Did not delete {rem_unr} .torrents(s) or content files.')
logger.dryrun(f'Did not delete {del_tor} .torrents(s) or content files.')
else:
logger.dryrun('No unregistered torrents found.')
else:
if rem_unr >= 1 or del_tor >= 1:
logger.info(f'Deleted {rem_unr} .torrents(s) but not content files.')
logger.info(f'Deleted {del_tor} .torrents(s) AND content files.')
else:
logger.info('No unregistered torrents found.')
def run():
update_category()
update_tags()
rem_unregistered()
cross_seed()
recheck()
if __name__ == '__main__':
run()