util.py 2.2 KB
Newer Older
W
wizardforcel 已提交
1 2 3 4 5 6 7 8 9 10
# -*- coding: utf-8 -*-

import os
import sys
from os import path
import hashlib
import types
import requests
import json
import time
W
wizardforcel 已提交
11
import tempfile
W
wizardforcel 已提交
12

W
wizardforcel 已提交
13
bundle_dir = tempfile.gettempdir()
W
wizardforcel 已提交
14 15 16

size_string = lambda byte: f"{byte / 1024 / 1024 / 1024:.2f} GB" if byte > 1024 * 1024 * 1024 else f"{byte / 1024 / 1024:.2f} MB" if byte > 1024 * 1024 else f"{byte / 1024:.2f} KB" if byte > 1024 else f"{int(byte)} B"

W
wizardforcel 已提交
17
def calc_sha1(data, hex=True):
W
wizardforcel 已提交
18
    sha1 = hashlib.sha1()
W
wizardforcel 已提交
19 20
    if hasattr(data, '__iter__') and \
       type(data) is not bytes:
W
wizardforcel 已提交
21 22 23 24
        for chunk in data:
            sha1.update(chunk)
    else:
        sha1.update(data)
W
wizardforcel 已提交
25
    return sha1.hexdigest() if hex else sha1.digest()
W
wizardforcel 已提交
26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53
    
    
def image_download(url):
    headers = {
        'Referer': "http://t.bilibili.com/",
        'User-Agent': "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/79.0.3945.79 Safari/537.36",
    }
    content = []
    last_chunk_time = None
    try:
        for chunk in requests.get(url, headers=headers, timeout=10, stream=True).iter_content(128 * 1024):
            if last_chunk_time is not None and time.time() - last_chunk_time > 5:
                return None
            content.append(chunk)
            last_chunk_time = time.time()
        return b"".join(content)
    except:
        return None
    

def read_history():
    try:
        with open(path.join(bundle_dir, "history.json"), "r", encoding="utf-8") as f:
            history = json.loads(f.read())
    except:
        history = {}
    return history

W
wizardforcel 已提交
54 55 56 57 58 59 60
def write_history(first_4mb_sha1, meta_dict, url):
    history = read_history()
    history[first_4mb_sha1] = meta_dict
    history[first_4mb_sha1]['url'] = url
    with open(os.path.join(bundle_dir, "history.json"), "w", encoding="utf-8") as f:
        f.write(json.dumps(history, ensure_ascii=False, indent=2))
    
W
wizardforcel 已提交
61 62 63
def read_in_chunk(fname, size=4 * 1024 * 1024, cnt=-1):
    with open(fname, "rb") as f:
        idx = 0
W
wizardforcel 已提交
64
        while True:
W
wizardforcel 已提交
65 66 67 68 69
            data = f.read(size)
            if not data or (cnt != -1 and idx >= cnt):
                break
            yield data
            idx += 1
W
wizardforcel 已提交
70 71 72
                
def log(message):
    print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] {message}")