pirate-get/pirate/print.py

import builtins
import re
import gzip
import urllib.parse as parse
import urllib.request as request
import shutil

import pirate.data

import colorama
import veryprettytable

from io import BytesIO
from http.cookiejar import CookieJar


class Printer:
    def __init__(self, enable_color):
        self.enable_color = enable_color

    def print(self, *args, **kwargs):
        if kwargs.get('color', False) and self.enable_color:
            colorama.init()
            color_dict = {
                'default': '',
                'header':  colorama.Back.BLACK + colorama.Fore.WHITE,
                'alt':     colorama.Fore.YELLOW,
                'zebra_0': '',
                'zebra_1': colorama.Fore.BLUE,
                'WARN':    colorama.Fore.MAGENTA,
                'ERROR':   colorama.Fore.RED}

            c = color_dict[kwargs.pop('color')]
            args = (c + args[0],) + args[1:] + (colorama.Style.RESET_ALL,)
            kwargs.pop('color', None)
            return builtins.print(*args, **kwargs)
        else:
            kwargs.pop('color', None)
            return builtins.print(*args, **kwargs)

    # TODO: extract the name from the search results
    #       instead of from the magnet link when possible
    def search_results(self, results, local=None):
        columns = shutil.get_terminal_size((80, 20)).columns
        even = True

        if local:
            table = veryprettytable.VeryPrettyTable(['LINK', 'DATE', 'SIZE', 'NAME'])

            table.align['SIZE'] = 'r'
            table.align['NAME'] = 'l'
        else:
            table = veryprettytable.VeryPrettyTable(['LINK', 'SEED', 'LEECH',
                                                     'RATIO', 'SIZE',
                                                     'UPLOAD', 'NAME'])
            table.align['NAME'] = 'l'
            table.align['SEED'] = 'r'
            table.align['LEECH'] = 'r'
            table.align['RATIO'] = 'r'
            table.align['SIZE'] = 'r'
            table.align['UPLOAD'] = 'l'

        table.max_width = columns
        table.border = False
        table.padding_width = 1

        for n, result in enumerate(results):

            name = re.search(r'dn=([^\&]*)', result['magnet'])
            torrent_name = parse.unquote_plus(name.group(1))

            if local:
                content = [n, result['date'], result['size'], torrent_name[:columns - 42]]
            else:
                no_seeders = int(result['seeds'])
                no_leechers = int(result['leechers'])
                if result['size'] != []:
                    size = float(result['size'][0])
                    unit = result['size'][1]
                else:
                    size = 0
                    unit = '???'
                date = result['uploaded']

                # compute the S/L ratio (Higher is better)
                try:
                    ratio = no_seeders / no_leechers
                except ZeroDivisionError:
                    ratio = float('inf')

                content = [n, no_seeders, no_leechers,
                           '{:.1f}'.format(ratio),
                           '{:.1f} '.format(size) + unit,
                           date, torrent_name[:columns - 50]]

            if even or not self.enable_color:
                table.add_row(content)
            else:
                table.add_row(content, fore_color='blue')

            # Alternate between colors
            even = not even
        self.print(table)

    def descriptions(self, chosen_links, results, site):
        jar = CookieJar()
        opener = request.build_opener(
            request.HTTPErrorProcessor,
            request.HTTPCookieProcessor(jar))

        for link in chosen_links:
            path = '/torrent/%s/' % results[link]['id']
            req = request.Request(site + path,
                                  headers=pirate.data.default_headers)
            req.add_header('Accept-encoding', 'gzip')
            f = opener.open(req, timeout=pirate.data.default_timeout)

            if f.info().get('Content-Encoding') == 'gzip':
                f = gzip.GzipFile(fileobj=BytesIO(f.read()))

            res = f.read().decode('utf-8')
            name = re.search(r'dn=([^\&]*)', results[link]['magnet'])
            torrent_name = parse.unquote(name.group(1)).replace('+', ' ')
            desc = re.search(r'<div class="nfo">\s*<pre>(.+?)(?=</pre>)',
                             res, re.DOTALL).group(1)

            # Replace HTML links with markdown style versions
            desc = re.sub(r'<a href="\s*([^"]+?)\s*"[^>]*>(\s*)([^<]+?)(\s*'
                          r')</a>', r'\2[\3](\1)\4', desc)

            self.print('Description for "%s":' % torrent_name, color='zebra_1')
            self.print(desc, color='zebra_0')

    def file_lists(self, chosen_links, results, site):
        jar = CookieJar()
        opener = request.build_opener(
            request.HTTPErrorProcessor,
            request.HTTPCookieProcessor(jar))

        for link in chosen_links:
            path = '/ajax_details_filelist.php'
            query = '?id=' + results[link]['id']
            req = request.Request(site + path + query,
                                  headers=pirate.data.default_headers)
            req.add_header('Accept-encoding', 'gzip')
            f = opener.open(req, timeout=pirate.data.default_timeout)

            if f.info().get('Content-Encoding') == 'gzip':
                f = gzip.GzipFile(fileobj=BytesIO(f.read()))

            # TODO: proper html decoding/parsing
            res = f.read().decode('utf-8').replace('&nbsp;', ' ')
            if 'File list not available.' in res:
                self.print('File list not available.')
                return
            files = re.findall(r'<td align="left">\s*([^<]+?)\s*</td><td ali'
                               r'gn="right">\s*([^<]+?)\s*</tr>', res)
            name = re.search(r'dn=([^\&]*)', results[link]['magnet'])
            torrent_name = parse.unquote(name.group(1)).replace('+', ' ')

            self.print('Files in "%s":' % torrent_name, color='zebra_1')
            cur_color = 'zebra_0'

            for f in files:
                self.print('{0[0]:>11}  {0[1]}'.format(f), color=cur_color)
                cur_color = 'zebra_0' if cur_color == 'zebra_1' else 'zebra_1'