156 lines
5.8 KiB
Python
156 lines
5.8 KiB
Python
import threading, isodate
|
|
import requests
|
|
import math
|
|
|
|
from requests.sessions import session
|
|
from tqdm import tqdm
|
|
from queue import Queue
|
|
|
|
dlthreads = 24
|
|
|
|
class WvDownloader(object):
|
|
def __init__(self, config):
|
|
self.mpd = config.mpd
|
|
self.output_file = config.output_file
|
|
self.mimetype = config.file_type
|
|
self.formatId = config.format_id
|
|
self.config = config
|
|
|
|
def process_url_templace(self, template, representation_id, bandwidth, time, number):
|
|
if representation_id is not None: result = template.replace('$RepresentationID$', representation_id)
|
|
if number is not None:
|
|
nstart = result.find('$Number')
|
|
if nstart >= 0:
|
|
nend = result.find('$', nstart+1)
|
|
if nend >= 0:
|
|
var = result[nstart+1 : nend]
|
|
if 'Number%' in var:
|
|
value = var[6:] % (int(number))
|
|
else:
|
|
value = number
|
|
result = result.replace('$'+var+'$', value)
|
|
if bandwidth is not None: result = result.replace('$Bandwidth$', bandwidth)
|
|
if time is not None: result = result.replace('$Time$', time)
|
|
result = result.replace('$$', '$').replace('../', '')
|
|
return result
|
|
|
|
def generate_segments(self):
|
|
segment_template = self.get_segment_template()
|
|
return self.get_segments(segment_template)
|
|
|
|
def get_segments(self, segment_template):
|
|
urls = []
|
|
urls.append(self.config.base_url + segment_template['@initialization'].replace('$RepresentationID$', self.config.format_id))
|
|
print(urls)
|
|
try:
|
|
current_number = int(segment_template.get("@startNumber", 0))
|
|
period_duration = self.get_duration()
|
|
segment_duration = int(segment_template["@duration"]) / int(segment_template["@timescale"])
|
|
total_segments = math.ceil(period_duration / segment_duration)
|
|
for _ in range(current_number, current_number + total_segments):
|
|
urls.append(self.config.base_url + self.process_url_templace(segment_template['@media'],
|
|
representation_id=self.config.format_id,
|
|
bandwidth=None, time="0", number=str(current_number)))
|
|
current_number += 1
|
|
except KeyError:
|
|
current_number = 0
|
|
current_time = 0
|
|
for seg in segment_template["SegmentTimeline"]["S"]:
|
|
if '@t' in seg:
|
|
current_time = seg['@t']
|
|
for i in range(int(seg.get('@r', 0)) + 1):
|
|
urls.append(self.config.base_url + self.process_url_templace(segment_template['@media'],
|
|
representation_id=self.config.format_id,
|
|
bandwidth=None, time=str(current_time), number=str(current_number)))
|
|
current_number += 1
|
|
current_time += seg['@d']
|
|
return urls
|
|
|
|
def get_duration(self):
|
|
media_duration = self.mpd["MPD"]["@mediaPresentationDuration"]
|
|
return isodate.parse_duration(media_duration).total_seconds()
|
|
|
|
def get_segment_template(self):
|
|
tracks = self.mpd['MPD']['Period']['AdaptationSet']
|
|
|
|
segment_template = []
|
|
if self.mimetype == "video/mp4":
|
|
for video_track in tracks:
|
|
if video_track["@mimeType"] == self.mimetype:
|
|
for v in video_track["Representation"]:
|
|
segment_template = v["SegmentTemplate"]
|
|
|
|
if self.mimetype == "audio/mp4":
|
|
for audio_track in tracks:
|
|
if audio_track["@mimeType"] == self.mimetype:
|
|
try:
|
|
segment_template = audio_track["SegmentTemplate"]
|
|
except (KeyError, TypeError):
|
|
for a in self.list_representation(audio_track):
|
|
segment_template = a["SegmentTemplate"]
|
|
|
|
return segment_template
|
|
|
|
def list_representation(self, x):
|
|
if isinstance(x['Representation'], list):
|
|
X = x['Representation']
|
|
else:
|
|
X = [x['Representation']]
|
|
return X
|
|
|
|
def run(self):
|
|
urls = self.generate_segments()
|
|
work_q = Queue()
|
|
result_q = Queue()
|
|
|
|
print('\n' + self.output_file)
|
|
pool = [WorkerThread(work_q=work_q, result_q=result_q) for i in range(dlthreads)]
|
|
for thread in pool:
|
|
thread.start()
|
|
|
|
work_count = 0
|
|
for seg_url in urls:
|
|
work_q.put((work_count, seg_url))
|
|
work_count += 1
|
|
results = []
|
|
|
|
for _ in tqdm(range(work_count)):
|
|
results.append(result_q.get())
|
|
outfile = open(self.output_file , 'wb+')
|
|
sortedr = sorted(results, key=lambda v: v[0])
|
|
for r in sortedr:
|
|
outfile.write(r[1])
|
|
outfile.close()
|
|
del results
|
|
print('Done!')
|
|
|
|
class Downloader:
|
|
def __init__(self):
|
|
self.session = requests.Session()
|
|
|
|
def DownloadSegment(self, url):
|
|
resp = self.session.get(url, stream=True)
|
|
resp.raw.decode_content = True
|
|
data = resp.raw.read()
|
|
return data
|
|
|
|
class WorkerThread(threading.Thread):
|
|
def __init__(self, work_q, result_q):
|
|
super(WorkerThread, self).__init__()
|
|
self.work_q = work_q
|
|
self.result_q = result_q
|
|
self.stoprequest = threading.Event()
|
|
self.downloader = Downloader()
|
|
|
|
def run(self):
|
|
while not self.stoprequest.isSet():
|
|
try:
|
|
(seq, url) = self.work_q.get(True, 0.05)
|
|
self.result_q.put((seq, self.downloader.DownloadSegment(url)))
|
|
except:
|
|
continue
|
|
|
|
def join(self, timeout=None):
|
|
self.stoprequest.set()
|
|
super(WorkerThread, self).join(timeout)
|