Compare commits

...

3 Commits

1 changed files with 187 additions and 170 deletions

View File

@ -1,7 +1,6 @@
#!/usr/bin/env python3 #!/usr/bin/env python3
import feedparser import feedparser
import subprocess
import click import click
import re import re
import json import json
@ -9,105 +8,114 @@ import hashlib
import sys import sys
import os import os
import transmissionrpc import transmissionrpc
from operator import itemgetter import requests
import base64
def printe(*args, **kwargs): def printe(*args, **kwargs):
print(*args, file=sys.stderr, **kwargs) print(*args, file=sys.stderr, **kwargs)
class Resolution: class Resolution:
# In decending order of quality, since indexes are used to tell whats "better" # In decending order of quality, since indexes are used to tell whats "better"
names = [ names = [
('8k', '8K', '4320p'), ('8k', '8K', '4320p'),
('5k', '5K', '2880p', 'UHD+'), ('5k', '5K', '2880p', 'UHD+'),
('4k', '4K', '2160p', 'UHD'), ('4k', '4K', '2160p', 'UHD'),
('2k', '2K', '1440p', 'QHD', 'WQHD'), ('2k', '2K', '1440p', 'QHD', 'WQHD'),
('1080p', 'FHD'), ('1080p', 'FHD'),
('720p', 'HD', 'HDTV'), ('720p', 'HD', 'HDTV'),
('480p', 'SD', 'VGA'), ('480p', 'SD', 'VGA'),
('360p', 'nHD'), ('360p', 'nHD'),
('240p', 'QVGA'), ('240p', 'QVGA'),
('Unknown', 'None', None) ('Unknown', 'None', None)
] ]
def __init__(self, _name):
self.name = _name
for family in self.names:
if self.name in family:
self.family = family
break
else:
raise ValueError("Did not recongnize resolution name. See Resolution.names")
def getNormalizedName(self): def __init__(self, _name):
return self.family[0] self.name = _name
for family in self.names:
if self.name in family:
self.family = family
break
else:
raise ValueError("Did not recongnize resolution name. See Resolution.names")
# Rich comparison support def getNormalizedName(self):
def __lt__(self, other): return self.family[0]
if not isinstance(other, Resolution):
raise TypeError # Rich comparison support
return self.names.index(self.family) > other.names.index(other.family) def __lt__(self, other):
def __eq__(self, other): if not isinstance(other, Resolution):
if not isinstance(other, Resolution): raise TypeError
raise TypeError return self.names.index(self.family) > other.names.index(other.family)
return self.names.index(self.family) == other.names.index(other.family)
def __ne__(self, other): def __eq__(self, other):
return not self.__eq__(other) if not isinstance(other, Resolution):
def __gt__(self, other): raise TypeError
return not self.__lt__(other) return self.names.index(self.family) == other.names.index(other.family)
def __le__(self, other):
return self.__lt__(other) or self.__eq__(other) def __ne__(self, other):
def __ge__(self, other): return not self.__eq__(other)
return self.__gt__(other) or self.__eq__(other)
def __gt__(self, other):
return not self.__lt__(other)
def __le__(self, other):
return self.__lt__(other) or self.__eq__(other)
def __ge__(self, other):
return self.__gt__(other) or self.__eq__(other)
class Show: class Show:
def __init__(self, name, episode=None): def __init__(self, name, episode=None):
self.name = name self.name = name
try: try:
self.episode = int(episode) if episode is not None else None self.episode = int(episode) if episode is not None else None
except (ValueError, TypeError): except (ValueError, TypeError):
raise ValueError("episode paramater must be a number or None") raise ValueError("episode paramater must be a number or None")
self.key = '{name} - {episode}'.format(name=self.name, episode=self.episode if self.episode else 'None') self.key = '{name} - {episode}'.format(name=self.name, episode=self.episode if self.episode else 'None')
def getConfig(self, config): def getConfig(self, config):
try: try:
show_config = next(want for want in config['wants'] if want['name'] == self.name) show_config = next(want for want in config['wants'] if want['name'] == self.name)
except StopIteration: except StopIteration:
show_config = None show_config = None
return show_config return show_config
def __repr__(self): def __repr__(self):
return "<Show {name} - {episode}>".format(name=self.name, episode=self.episode if self.episode else 'None') return "<Show {name} - {episode}>".format(name=self.name, episode=self.episode if self.episode else 'None')
def __hash__(self): def __hash__(self):
return hash(self.key) return hash(self.key)
def __eq__(self, other):
return self.key == other.key
def __eq__(self, other):
return self.key == other.key
class Torrent: class Torrent:
def __init__(self, name, link, id=None, episode=None, quality=None, encoding=None): def __init__(self, name, link, id=None, episode=None, quality=None, encoding=None):
self.name = name self.name = name
self.link = link self.link = link
self.id = id if id else hashlib.sha256(link).hexdigest() self.id = id if id else hashlib.sha256(link).hexdigest()
try: try:
self.episode = int(episode) if episode is not None else None self.episode = int(episode) if episode is not None else None
except (ValueError, TypeError): except (ValueError, TypeError):
raise ValueError("episode paramater must be a number or None") raise ValueError("episode paramater must be a number or None")
self.quality = Resolution(quality) self.quality = Resolution(quality)
self.encoding = encoding self.encoding = encoding
def isShow(self, show): def isShow(self, show):
return self.name == show.name and (show.episode is None or self.episode == show.episode) return self.name == show.name and (show.episode is None or self.episode == show.episode)
def __repr__(self): def __repr__(self):
return "<Torrent {id} for {name} - {episode}>".format(id=self.id, name=self.name, episode=self.episode) return "<Torrent {id} for {name} - {episode}>".format(id=self.id, name=self.name, episode=self.episode)
def __hash__(self): def __hash__(self):
return hash(self.id) return hash(self.id)
def __eq__(self, other):
return self.id == other.id
def __eq__(self, other):
return self.id == other.id
@click.command() @click.command()
@ -116,102 +124,111 @@ class Torrent:
@click.option('--dry-run', 'dry_run', type=click.BOOL, is_flag=True, default=False, help="Don't act on any torrents found. Implies --no-dedup") @click.option('--dry-run', 'dry_run', type=click.BOOL, is_flag=True, default=False, help="Don't act on any torrents found. Implies --no-dedup")
@click.option('--no-dedup', 'no_seen', type=click.BOOL, is_flag=True, default=False, help="Don't save seen torrents to the seen file.") @click.option('--no-dedup', 'no_seen', type=click.BOOL, is_flag=True, default=False, help="Don't save seen torrents to the seen file.")
def main(config_path='config.json', seen_path='seen.txt', dry_run=False, no_seen=False): def main(config_path='config.json', seen_path='seen.txt', dry_run=False, no_seen=False):
# Get our config # Get our config
with open(config_path, 'r') as config_file: with open(config_path, 'r') as config_file:
config_json = ''.join(config_file.readlines()) config_json = ''.join(config_file.readlines())
config = json.loads(config_json) config = json.loads(config_json)
providers = config['providers'] providers = config['providers']
wants = [want['name'] for want in config['wants']] wants = [want['name'] for want in config['wants']]
seen = [line.strip() for line in open(seen_path, 'r')] seen = [line.strip() for line in open(seen_path, 'r')]
filtered = {} filtered = {}
for provider in providers: for provider in providers:
# Compile and check our regex # Compile and check our regex
title_parser = re.compile(provider['regex']) title_parser = re.compile(provider['regex'])
if title_parser.groupindex.get('name') is None: if title_parser.groupindex.get('name') is None:
printe("ERROR: Regex for {provider_name} does not contain a 'name' group. Name a capture group in your regex using '(?P<name>...)'".format(provider_name=provider['name'])) printe("ERROR: Regex for {provider_name} does not contain a 'name' group. Name a capture group in your regex using '(?P<name>...)'".format(provider_name=provider['name']))
continue continue
elif title_parser.groupindex.get('quality') is None or title_parser.groupindex.get('episode') is None: elif title_parser.groupindex.get('quality') is None or title_parser.groupindex.get('episode') is None:
printe("WARNING: Regex for {provider_name} does not contain an 'episode' and/or 'quality' group. Unless you are sure duplicates don't exist in the feeds, add optional groups using '(?P<name>...)?'".format(provider_name=provider['name'])) printe("WARNING: Regex for {provider_name} does not contain an 'episode' and/or 'quality' group. Unless you are sure duplicates don't exist in the feeds, add optional groups using '(?P<name>...)?'".format(provider_name=provider['name']))
# Get and parse rss feed # Get and parse rss feed
rss = feedparser.parse(provider['url']) rss = feedparser.parse(provider['url'])
for entry in rss.entries: for entry in rss.entries:
# Run regex against rss entry title (usually the file / torrent name) # Run regex against rss entry title (usually the file / torrent name)
matched_title = title_parser.match(entry.title) matched_title = title_parser.match(entry.title)
if matched_title is None: if matched_title is None:
continue continue
parsed_title = matched_title.groupdict() parsed_title = matched_title.groupdict()
torrent = Torrent( torrent = Torrent(
parsed_title['name'], parsed_title['name'],
entry.get('link') or entry.links[0].href, entry.get('link') or entry.links[0].href,
id=entry.get('id'), id=entry.get('id'),
episode=parsed_title.get('episode'), episode=parsed_title.get('episode'),
quality=parsed_title.get('quality'), quality=parsed_title.get('quality'),
encoding=parsed_title.get('encoding') encoding=parsed_title.get('encoding')
) )
show = Show(torrent.name, episode=torrent.episode) show = Show(torrent.name, episode=torrent.episode)
# Check if its something we're interested in and add it to a dict # Check if its something we're interested in and add it to a dict
if torrent.name in wants and torrent.id not in seen: if torrent.name in wants and torrent.id not in seen:
printe("Found new torrent: {torrent}".format(torrent=entry.title)) printe("Found new torrent: {torrent}".format(torrent=entry.title))
if filtered.get(show) is None: if filtered.get(show) is None:
filtered[show] = [torrent] filtered[show] = [torrent]
else: else:
filtered[show].append(torrent) filtered[show].append(torrent)
# Record that we have seen it # Record that we have seen it
seen.append(torrent.id) seen.append(torrent.id)
# Configure a transmissionrpc client if set # Configure a transmissionrpc client if set
if config.get('rpc'): if config.get('rpc'):
# Allows for a setting like `"rpc": true` to just use defaults # Allows for a setting like `"rpc": true` to just use defaults
# HACK: Should just do the config['rpc'].get better # HACK: Should just do the config['rpc'].get better
if isinstance(config['rpc'], bool): if isinstance(config['rpc'], bool):
config['rpc'] = {} config['rpc'] = {}
rpc = transmissionrpc.Client( rpc = transmissionrpc.Client(
address=config['rpc'].get('address', 'localhost'), address=config['rpc'].get('address', 'localhost'),
port=config['rpc'].get('port', 9091), port=config['rpc'].get('port', 9091),
user=config['rpc'].get('user'), user=config['rpc'].get('user'),
password=config['rpc'].get('password'), password=config['rpc'].get('password'),
http_handler=config['rpc'].get('http_handler'), http_handler=config['rpc'].get('http_handler'),
timeout=config['rpc'].get('timeout') timeout=config['rpc'].get('timeout')
) )
default_download_dir = rpc.get_session().download_dir default_download_dir = rpc.get_session().download_dir
# Choose the best item else:
for show, torrents in filtered.items(): rpc = None
item_config = show.getConfig(config) # Choose the best item
# Filter for show, torrents in filtered.items():
if item_config.get('target-quality'): item_config = show.getConfig(config)
torrents = [torrent for torrent in torrents if torrent.quality == Resolution(item_config['target-quality'])] # Filter
elif item_config.get('min-quality'): if item_config.get('target-quality'):
torrents = [torrent for torrent in torrents if torrent.quality >= Resolution(item_config['min-quality'])] torrents = [torrent for torrent in torrents if torrent.quality == Resolution(item_config['target-quality'])]
elif item_config.get('min-quality'):
torrents = [torrent for torrent in torrents if torrent.quality >= Resolution(item_config['min-quality'])]
if len(torrents) > 0: if len(torrents) > 0:
best_entry = sorted(torrents, key=lambda t: t.quality, reverse=True)[0] best_entry = sorted(torrents, key=lambda t: t.quality, reverse=True)[0]
# Output the magnet / torrent link # Output the magnet / torrent link
printe("Choosing {torrent} as best choice for {show} (quality: {res})".format(torrent=best_entry.id, show=show.name, res=best_entry.quality.getNormalizedName())) printe("Choosing {torrent} as best choice for {show} (quality: {res})".format(torrent=best_entry.id, show=show.name, res=best_entry.quality.getNormalizedName()))
# TODO: transmission-cli subprocess call here? # TODO: transmission-cli subprocess call here?
print(best_entry.link) print(best_entry.link)
# Add torrent to transmission directly # Add torrent to transmission directly
if rpc and not dry_run: if rpc and not dry_run:
destination = item_config.get('destination') destination = item_config.get('destination')
if destination: if destination:
destination = destination.format(show=show, torrent=torrent) destination = destination.format(show=show, torrent=torrent)
destination = os.path.expanduser(destination) destination = os.path.expanduser(destination)
if not destination.startswith('/'): if not destination.startswith('/'):
destination = os.path.join(default_download_dir, destination) destination = os.path.join(default_download_dir, destination)
if not os.path.exists(destination): if not os.path.exists(destination):
os.makedirs(destination) os.makedirs(destination)
t = rpc.add_torrent(best_entry.link, download_dir=destination or default_download_dir)
else: if best_entry.link.startswith('http'):
printe("No torrents for {show} pass quality restrictions".format(show=show.name)) r = requests.get(best_entry.link)
r.raise_for_status()
rpc.add_torrent(base64.b64encode(r.content).decode('utf-8'), download_dir=destination or default_download_dir)
else:
rpc.add_torrent(best_entry.link, download_dir=destination or default_download_dir)
else:
printe("No torrents for {show} pass quality restrictions".format(show=show.name))
# Save our records to a file for later use
if not (dry_run or no_seen):
with open(seen_path, 'w') as seen_file:
seen_file.write('\n'.join(seen))
# Save our records to a file for later use
if not (dry_run or no_seen):
with open(seen_path, 'w') as seen_file:
seen_file.write('\n'.join(seen))
# Usual python stuff # Usual python stuff
if __name__ == '__main__': if __name__ == '__main__':
main() main()