# Скрипт выводит в консоль содержимое файла в двоичном виде
# Можно использовать перенаправление вывода для сохранения файла
# Использование vkontakte_m3u8_downloader.py ссылка на аудиозапись > имя файла
# Для работы скрипта нужен ffmpeg
import re
import os.path
import sys
import tempfile
import subprocess
import urllib
from urllib.request import urlopen
from Crypto.Cipher import AES
from Crypto.Util.Padding import unpad

import requests

regex = "https.*key.pub"
regex_2 = '"https.*pub?.*"'
regex_3 = 'ts?.*'


def parse_m3u8(_url):
    urldir = os.path.dirname(_url)
    playlist = requests.get(_url).text

    # парсим секции #EXT-X-KEY, складываем в список [файл, метод шифрования, url к ключу]
    chunks = []
    for match in re.finditer(r'(#EXT-X-KEY:(?:(?!#EXT-X-KEY:).)*)', playlist, re.DOTALL):
        # получили секцию, смотрим метод шифрования и URL к файлу ключа
        section = match.group(0)
        section_match = re.match(r'.*?METHOD=(?P<METHOD>[^,\n\r]+)(?:,URI=(?P<URI>[^\n]+)|)', section, re.IGNORECASE)
        key_method = section_match.group('METHOD')
        key_url = section_match.group('URI')
        if key_url:
            key_url = key_url.strip("'\"")
        # смотрим список файлов в этой секции, создаем список
        for match2 in re.finditer(r'.*?(?P<file>^[^\r\n]+\.ts)(?P<extra>[^\n\r]+)', section, re.IGNORECASE | re.DOTALL | re.MULTILINE):
            file = match2.group('file')
            extra = match2.group('extra')
            file_url = f'{urldir}/{file}{extra}'
            chunks.append({'file_url': file_url, 'key_method': key_method, 'key': key_url, 'file_name': file})

    # для созданного списка файлов, достаем содержимое ключей, используя кэширование
    cached_urls = []
    for chunk in chunks:
        key_url = chunk['key']
        if not key_url:
            continue
        value = ''

        # если url ключа уже скачивалась, достаём кэшированное значение
        for cached_url in cached_urls:
            if key_url == cached_url[0]:
                value = cached_url[1]
                break
        # делаем запрос только если значение не найдено, и кэшируем ответ
        if not value:
            value = requests.get(key_url).text
            cached_urls.append([key_url, value])
        # в общем списке заменяем url ключа на его значение
        chunk['key'] = value
    del cached_urls
    return [chunks, playlist]


def download_m3u8(job, output='out.ts'):
    chunks = job[0]
    playlist = job[1]

    with tempfile.TemporaryDirectory() as _dir:
        os.chdir(_dir)
        print(f'Скачиваем части в каталог {_dir}')

        file_num = 1
        for chunk in chunks:
            file_name = str(file_num).zfill(4) + " " + chunk['file_name']
            file_num += 1

            if chunk['key']:
                key = chunk['key']
                key = key.encode()
                file_in = urlopen(chunk['file_url'])
                iv = file_in.read(16) # Read the iv out - this is 16 bytes long
                ciphered_data = file_in.read() # Read the rest of the data
                file_in.close()
                cipher = AES.new(key, AES.MODE_CBC, iv=iv)  # Setup cipher
                data = unpad(cipher.decrypt(ciphered_data), AES.block_size) # Decrypt and then up-pad the result
            else:
                data = requests.get(chunk['file_url']).content

            if data:
                with open(file_name, 'wb') as audio:
                    audio.write(data)
                    path = f'{_dir}\\{file_name}'
                    chunk['file_url'] = path
                    print(f'Скачан файл {file_name}')
            else:
                print(f'Ошибка получения части {chunk["file_name"]}')

        # склеиваем файлы
        ffmpeg_queue = '# queue'
        for chunk in chunks:
            ffmpeg_queue += f"\nfile '{chunk['file_url']}'"

        queue_lst = 'queue.txt'
        with open(queue_lst, 'w') as file:
            file.write(ffmpeg_queue)

        with open('playlist.txt', 'w') as file:
            file.write(playlist)

        subprocess.call(['ffmpeg', '-f', 'concat', '-safe', '0', '-i', f'{queue_lst}', '-c', 'copy', f'{output}'])
    return chunks



if __name__ == '__main__':
    url = sys.argv[1]
    result = parse_m3u8(url)
    sys.stdout.buffer.write(download_m3u8(result))