python下载m3u8视频文件优化3

优化

ts优化：

兼容
1.m3u8文件中的ts地址完整
2.m3u8文件中的ts地址不完整，需要输入ts地址的前部分

解密优化

兼容
1.无需解密ts
2.需解密，但需要手动输入key_url
3.需解密，但不需要手动输入key_url

执行过程

def main():
    if not os.path.isdir('video'):
        os.mkdir('video')
    title = input('请输入电影名称：').strip()
    # 1 是否需要修改ts地址？
    ts = input('是否需要修改ts地址（Y/N):').strip().upper()
    # 1-1 如果需要修改，手动输入ts地址前部
    if ts == 'Y':
        ts_url_head = input('请输入需要拼接的ts地址：').strip()
        lst, lst_n = get_lst(ts_url_head)
     # 1-2 如果不需要修改，就表明m3u8文件中的ts地址是完整的
    else:
        lst, lst_n = get_lst()
    thread_download(lst, lst_n)
    # 2 是否需要手动输入key_url
    manual = input('是否需要手动输入key_url（Y/N）：').strip().upper()
    # 2-1 如果需要手动输入key地址，一定是要解密的，将key地址传入代码
    if manual == 'Y':
        key_url = input('请输入key_url：').strip()
        thread_decode(lst, lst_n, key_url)
    # 2-2 如果不需要修改，分为两种情况，一是不需要解密，另一种情况是需要解密，但可以通过re从m3u8文件中提取完整的key_url.
    else:
        key = judgment()
        if key is not None:
            key_url = key.group(1)
            thread_decode(lst, lst_n, key_url)
    merge(len(lst), title)

完整代码

import os
from concurrent.futures import ThreadPoolExecutor
import re
import requests
from Crypto.Cipher import AES
from tqdm import tqdm


def get_lst(ts_url_head=''):
    task = []
    num = []
    with open('1.m3u8', 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if line.startswith('#'):
                continue
            url = ts_url_head + line
            task.append(url)
    for i in range(len(task)):
        num.append(i)
    return task, num


def download(url, name):
    resp = requests.get(url, timeout=10)
    with open(f'video/{name}', 'wb') as f1:
        f1.write(resp.content)
        print(f'{name}下载完成')


def thread_download(lst, lst_n):
    with ThreadPoolExecutor(50) as t:
        for i in lst_n:
            t.submit(download, name=i, url=lst[i])
    check(lst, lst_n)


def check(lst, lst_n):
    lst_ns = []
    for i in lst_n:
        if not os.path.isfile(f'video/{i}'):
            lst_ns.append(i)
    if lst_ns:
        print(lst_ns)
        print('开始递归==================================================')
        thread_download(lst, lst_ns)
    else:
        print('全部下载完成')


def merge(num, title):
    with open(f'video/{title}.mp4', 'wb') as f:
        for i in tqdm(range(num)):
            if os.path.isfile(f'video/{i}_1'):
                with open(f'video/{i}_1', 'rb') as f1:
                    f.write(f1.read())
            else:
                with open(f'video/{i}', 'rb') as f1:
                    f.write(f1.read())
    print('合并完成')
    delete_ts(num)


def decode_ts(name, url):
    resp = requests.get(url)
    key = resp.content
    iv = b'0000000000000000'
    aes = AES.new(key, AES.MODE_CBC, iv)
    with open(f'video/{name}', 'rb') as f:
        de_f = aes.decrypt(f.read())
        with open(f'video/{name}_1', 'wb') as f1:
            f1.write(de_f)
    print(f'{name}解码完成')


def thread_decode(lst, lst_n, key_url):
    with ThreadPoolExecutor(50) as t:
        for i in lst_n:
            t.submit(decode_ts, name=i, url=key_url)
    check_de_ts(lst, lst_n, key_url)


def check_de_ts(lst, lst_n, key_url):
    lst_ns = []
    for i in lst_n:
        if not os.path.isfile(f'video/{i}_1'):
            lst_ns.append(i)
    if lst_ns:
        print(lst_ns)
        print('开始递归==================================================')
        thread_download(lst, lst_ns)
        thread_decode(lst, lst_ns, key_url)
    else:
        print('全部ts解密完成')


def delete_ts(num):
    for i in range(num):
        os.remove(f'video/{i}')
        if os.path.isfile(f'video/{i}_1'):
            os.remove(f'video/{i}_1')


def judgment():
    with open('1.m3u8', 'r', encoding='utf-8') as f:
        obj = re.compile(r'URI="(.*?)"')
        key = obj.search(f.read())
    return key


def main():
    if not os.path.isdir('video'):
        os.mkdir('video')
    title = input('请输入电影名称：').strip()
    ts = input('是否需要修改ts地址（Y/N):').strip().upper()
    if ts == 'Y':
        ts_url_head = input('请输入需要拼接的ts地址：').strip()
        lst, lst_n = get_lst(ts_url_head)
    else:
        lst, lst_n = get_lst()
    thread_download(lst, lst_n)
    manual = input('是否需要手动输入key_url（Y/N）：').strip().upper()
    if manual == 'Y':
        key_url = input('请输入key_url：').strip()
        thread_decode(lst, lst_n, key_url)
    else:
        key = judgment()
        if key is not None:
            key_url = key.group(1)
            thread_decode(lst, lst_n, key_url)
    merge(len(lst), title)


if __name__ == '__main__':
    main()