CDN文件中转

file_transfer.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Author: Ficapy

from __future__ import division, print_function
import sys
import re
import urllib
import sqlite3
import tempfile
import time
import string
import os
import json
import wget
import requests
import youtube_dl
from math import ceil
from subprocess import check_output
from hashlib import md5
from base64 import b64encode
from requests_toolbelt import MultipartEncoder, MultipartEncoderMonitor
from progressbar import ProgressBar, Percentage, Bar, ETA, FileTransferSpeed

SEC = '-------------'
NAME = '------------'
PWD = '--------------'
url = 'http://m0.api.upyun.com/----/'
bucket = '-----------'


def sign(params, token_secret=None):
    _sec = token_secret or SEC
    signature = md5(''.join(key + str(value) for key, value in sorted(params.items())) + _sec).hexdigest()
    policy = b64encode(json.dumps(params))
    return {'signature': signature,
            'policy': policy}


def safe_file_name(str):
    valid_chars = "-_.() %s%s" % (string.ascii_letters, string.digits)
    return ''.join(c for c in str if c in valid_chars)


class ProgressBarHandler(object):
    def __init__(self, params, totalsize):
        widgets = [params, ' ', Percentage(), ' ',
                   Bar(marker='=', left='[', right=']'), ' ',
                   ETA(), ' ', FileTransferSpeed()]
        self.pbar = ProgressBar(widgets=widgets, maxval=totalsize).start()

    def update(self, readsofar):
        try:
            self.pbar.update(readsofar + self.pbar.currval)
        except ValueError:
            pass

    def show(self, value):
        try:
            self.pbar.update(value)
        except ValueError:
            pass

    def finish(self):
        self.pbar.finish()


# DOWNLOAD

def cache(func):
    def _cache(url):
        conn = sqlite3.connect('{}/file_download_record.db'.format(tempfile.gettempdir()))
        c = conn.cursor()
        c.execute("""CREATE TABLE IF NOT EXISTS download (url TEXT, time TEXT,file_path TEXT)""")
        c.execute('SELECT * FROM download WHERE url=?', (url,))
        exists = c.fetchall()
        if exists:
            for i in exists:
                if int(i[-2]) + 1800 > int(time.time()):
                    print('file exists,ignore download process')
                    c.close()
                    print(i[-1])
                    return i[-1]
        file_path = func(url)
        c.execute("INSERT INTO download VALUES (?,?,?)", (url, int(time.time()), file_path))
        conn.commit()
        c.close()
        return file_path

    return _cache


@cache
def youtube_download(url):
    class MyLogger(object):
        def debug(self, msg):
            pass

        def warning(self, msg):
            pass

        def error(self, msg):
            print(msg)

    def my_hook(d):
        pbar = ProgressBarHandler('  Download', d['total_bytes'])
        if d['status'] == 'finished':
            filename = d['filename']
            tf = tempfile.gettempdir()
            dst = os.path.join(tf, safe_file_name(filename))
            my_hook.ret = dst
            os.rename(filename, dst)
            pbar.finish()
        else:
            pbar.show(d['downloaded_bytes'])

    my_hook.ret = ''

    ydl_opts = {
        'logger': MyLogger(),
        'progress_hooks': [my_hook],
    }
    with youtube_dl.YoutubeDL(ydl_opts) as ydl:
        ydl.download([url])

    return my_hook.ret


@cache
def normal_download(url):
    def wbar(current_size, total_size, width, _cache=[]):
        if _cache:
            bar = _cache[0]
        else:
            bar = ProgressBarHandler('  Download', total_size)
            _cache.append(bar)
        if abs(current_size - total_size) < 100:
            bar.finish()
        else:
            bar.show(current_size)

    path = tempfile.gettempdir()

    filename = wget.download(url, out=path, bar=wbar)

    file_path = os.path.join(path, filename)
    dst = os.path.join(path, safe_file_name(filename))
    os.rename(file_path, dst)

    return dst


def download(url):
    if 'youtube' in url:
        return youtube_download(url)
    return normal_download(url)


# UPLOAD

def init(path):
    assert os.path.isfile(path)
    file_blocks = int(ceil(os.path.getsize(path) / 4 / 1024 / 1024))
    file_size = os.path.getsize(path)
    params = {
        'path': os.path.split(path)[-1],
        'expiration': int(time.time()) + 24 * 3600,
        'file_blocks': file_blocks,
        'file_hash': check_output(['openssl', 'md5', path]).split()[-1],
        'file_size': file_size
    }
    ret = requests.post(url, data=sign(params), auth=(NAME, PWD)).json()
    upload_bar = ProgressBarHandler(params=' Uploading', totalsize=file_size)
    return {
        'save_token': ret['save_token'],
        'expiration': ret['expired_at'],
        'token_secret': ret['token_secret'],
        'blocks': file_blocks,
        'upload_bar': upload_bar,
        'status': ret['status']
    }


def chuck_upload(path, block_index, save_token, expiration, token_secret, **kwargs):
    # every single chuck size is 4MB
    if init(path)['status'][block_index] == 1:
        return
    tf = tempfile.TemporaryFile()
    with open(path, 'rb') as r, open(tf.name, 'w+') as w:
        r.seek(block_index * 1024 * 1024 * 4)
        w.write(r.read(1024 * 1024 * 4))
    hash = md5(open(tf.name).read()).hexdigest()
    params = {
        'save_token': save_token,
        'expiration': expiration,
        'block_index': block_index,
        'block_hash': hash
    }
    params = sign(params, token_secret=token_secret)
    params.update({'file': ('filename', open(tf.name, 'rb'))})

    def create_callback(encoder):
        def callback(monitor):
            kwargs['upload_bar'].show(monitor.bytes_read + block_index * 1024 * 1024 * 10)

        return callback

    encoder = MultipartEncoder(params)
    callback = create_callback(encoder)
    monitor = MultipartEncoderMonitor(encoder, callback)
    requests.post(url, data=monitor, headers={'Content-Type': monitor.content_type})

    os.remove(tf.name)
    return


def merge(save_token, expiration, token_secret, **kwargs):
    params = {
        'save_token': save_token,
        'expiration': expiration
    }
    ret = requests.post(url, data=sign(params, token_secret=token_secret))
    path = ret.json()['path']
    print('http://{}.b0.upaiyun.com{}'.format(bucket, urllib.quote(path)))


def main(url):
    path = download(url)
    setp_1 = init(path)
    for i in range(setp_1['blocks']):
        chuck_upload(path, i, **setp_1)
    setp_1['upload_bar'].finish()
    merge(**setp_1)


if __name__ == '__main__':
    url_ = sys.argv[-1]
    if not re.findall('http[s]?://(?:[a-zA-Z]|[0-9]|[$-_@.&+]|[!*\(\),]|(?:%[0-9a-fA-F][0-9a-fA-F]))+', url_):
        raise Exception('must be a full url link')
    main(url_)

https://asciinema.org/a/28559