util.py

#!/usr/bin/env python2
# coding=utf-8
from __future__ import print_function

import base64
import urlparse
import logging
import pickle
from requests.utils import dict_from_cookiejar

from command.config import global_config

import bddown_help

__all__ = [
    "bd_help",
    "usage",
    "parse_url",
    "add_http",
    "convert_none",
    "bcolor",
    "in_list",
    "logger",
]

URL = ['pan.baidu.com', 'yun.baidu.com']
FILTER_KEYS = ['shareid', 'server_filename', 'isdir', 'fs_id', 'sign', 'time_stamp', 'shorturl', 'dlink',
               'filelist', 'operation']
# TODO: add md5


def bd_help(args):
    if len(args) == 1:
        helper = getattr(bddown_help, args[0].lower(), bddown_help.help)
        usage(helper)
    elif len(args) == 0:
        usage(bddown_help.show_help)
    else:
        usage(bddown_help.help)


def usage(doc=bddown_help.usage, message=None):
    if hasattr(doc, '__call__'):
        doc = doc()
    if message:
        print(message)
    print(doc.strip())


def parse_url(url):
    """
    This function will parse url and judge which type the link is.

    :type url: str
    :param url: baidu netdisk share url.
    :return: dict
    """
    result = urlparse.urlparse(url)

    # wrong url
    if result.netloc not in ('pan.baidu.com', 'yun.baidu.com'):
        return {'type': -1}

    # http://pan.baidu.com/s/1kTFQbIn or http://pan.baidu.com/share/link?shareid=2009678541&uk=2839544145
    if result.path.startswith('/s/') or ('link' in result.path):
        return {'type': 1}

    # http://pan.baidu.com/share/verify?shareid=2009678541&uk=2839544145
    elif 'init' in result.path:
        return {'type': 1}

    # FIXME: Url could be the album type
    # eg: http://pan.baidu.com/wap/album/info?uk=2166491526&album_id=4852578710285570610&third=0
    # and http://pan.baidu.com/wap/album/file?uk=2166491526&album_id=4852578710285570610&fsid=1086862507948619
    # http://pan.baidu.com/pcloud/album/info?uk=3943531277&album_id=1553987381796453514
    elif 'album' in result.path:
        info = dict(urlparse.parse_qsl(result.query))
        info['type'] = 2
        return info

    # TODO: download share home
    # http://pan.baidu.com/share/home?uk=NUMBER
    elif 'home' in result.path and result.query:
        return {'type': 3}

    # Fix #17
    # Workaround: Redirect wap page to PC page
    elif 'wap' in result.path and 'fsid' in result.query:
        params = urlparse.parse_qs(result.query)
        fs_id = params.get('fsid')
        share_id = params.get('shareid')
        uk = params.get('uk')
        if not fs_id or not share_id or not uk:
            return {'type': 2}
        url = 'http://pan.baidu.com/share/link?uk={uk}&shareid={shareid}'.format(uk=uk[0], shareid=share_id[0])
        return {'type': 4, 'fsid': fs_id[0], 'url': url}

    else:
        return {'type': 0}


def add_http(url):
    if url.startswith('http://') or url.startswith('https://'):
        return url
    else:
        return 'http://' + url

convert_none = lambda opt, arg: opt + arg if arg else ""


# from http://stackoverflow.com/questions/287871/print-in-terminal-with-colors-using-python
# THANKS!

class BColor(object):
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'

    def disable(self):
        self.HEADER = ''
        self.OKBLUE = ''
        self.OKGREEN = ''
        self.WARNING = ''
        self.FAIL = ''
        self.ENDC = ''


bcolor = BColor()

in_list = lambda key, want_keys: key in want_keys


def filter_dict(bool_func, dictionary, want_keys):
    filtered_dict = {}
    for each_key in dictionary.keys():
        if bool_func(each_key, want_keys):
            filtered_dict[each_key] = dictionary[each_key]
    return filtered_dict


def merge_dict(dictionary, key):
    # will remove
    try:
        dictionary.update(dictionary[key][0])
        del dictionary[key]
    except KeyError:
        pass
    return dictionary


def filter_dict_wrapper(dictionary):
    d = {}
    for (k, v) in dictionary.items():
        if k in FILTER_KEYS:
            d[k] = v
        elif k == 'filelist':
            d[k] = [filter_dict(in_list, item, FILTER_KEYS) for item in v]
        elif k == 'operation':
            d[k] = [filter_dict(in_list, item, FILTER_KEYS) for item in v[0].get('filelist')]
    return d


def hack_sign(sign3, sign1):
    """
    Generate sign which is needed by downloading private file.
    Hack from `yunData.sign2`.

    :param sign3: yunData.sign3
    :type sign3: str
    :param sign1: yunData.sign1
    :type sign1: str
    :return: str (base64 encoded string)
    """
    def sign2(s3, s1):
        o = ""
        v = len(s3)
        a = [ord(s3[i % v]) for i in range(256)]
        p = range(256)
        # loop one
        u = 0
        for q in range(256):
            u = (u + p[q] + a[q]) % 256
            p[q], p[u] = p[u], p[q]
        # loop two
        i = u = 0
        for q in range(len(s1)):
            i = (i + 1) % 256
            u = (u + p[i]) % 256
            p[i], p[u] = p[u], p[i]
            k = p[((p[i] + p[u]) % 256)]
            o += chr(ord(s1[q]) ^ k)
        return o
    return base64.encodestring(sign2(sign3, sign1)).rstrip("\n")


def get_logger(logger_name):
    alogger = logging.getLogger(logger_name)
    fmt = logging.Formatter("%(levelname)s - %(method)s - %(type)s: \n-> %(message)s\n")
    handler = logging.StreamHandler()
    handler.setFormatter(fmt)
    alogger.addHandler(handler)
    alogger.setLevel(logging.INFO)
    return alogger


def save_cookies(cookies):
    with open(global_config.cookies, 'w') as f:
        pickle.dump(dict_from_cookiejar(cookies), f)


logger = get_logger('pan')