Merge pull request #5588 from aajanki/encode_frag_filenames
[youtube-dl] / youtube_dl / downloader / f4m.py
1 from __future__ import division, unicode_literals
2
3 import base64
4 import io
5 import itertools
6 import os
7 import time
8 import xml.etree.ElementTree as etree
9
10 from .fragment import FragmentFD
11 from ..compat import (
12     compat_urlparse,
13     compat_urllib_error,
14 )
15 from ..utils import (
16     struct_pack,
17     struct_unpack,
18     xpath_text,
19 )
20
21
22 class FlvReader(io.BytesIO):
23     """
24     Reader for Flv files
25     The file format is documented in https://www.adobe.com/devnet/f4v.html
26     """
27
28     # Utility functions for reading numbers and strings
29     def read_unsigned_long_long(self):
30         return struct_unpack('!Q', self.read(8))[0]
31
32     def read_unsigned_int(self):
33         return struct_unpack('!I', self.read(4))[0]
34
35     def read_unsigned_char(self):
36         return struct_unpack('!B', self.read(1))[0]
37
38     def read_string(self):
39         res = b''
40         while True:
41             char = self.read(1)
42             if char == b'\x00':
43                 break
44             res += char
45         return res
46
47     def read_box_info(self):
48         """
49         Read a box and return the info as a tuple: (box_size, box_type, box_data)
50         """
51         real_size = size = self.read_unsigned_int()
52         box_type = self.read(4)
53         header_end = 8
54         if size == 1:
55             real_size = self.read_unsigned_long_long()
56             header_end = 16
57         return real_size, box_type, self.read(real_size - header_end)
58
59     def read_asrt(self):
60         # version
61         self.read_unsigned_char()
62         # flags
63         self.read(3)
64         quality_entry_count = self.read_unsigned_char()
65         # QualityEntryCount
66         for i in range(quality_entry_count):
67             self.read_string()
68
69         segment_run_count = self.read_unsigned_int()
70         segments = []
71         for i in range(segment_run_count):
72             first_segment = self.read_unsigned_int()
73             fragments_per_segment = self.read_unsigned_int()
74             segments.append((first_segment, fragments_per_segment))
75
76         return {
77             'segment_run': segments,
78         }
79
80     def read_afrt(self):
81         # version
82         self.read_unsigned_char()
83         # flags
84         self.read(3)
85         # time scale
86         self.read_unsigned_int()
87
88         quality_entry_count = self.read_unsigned_char()
89         # QualitySegmentUrlModifiers
90         for i in range(quality_entry_count):
91             self.read_string()
92
93         fragments_count = self.read_unsigned_int()
94         fragments = []
95         for i in range(fragments_count):
96             first = self.read_unsigned_int()
97             first_ts = self.read_unsigned_long_long()
98             duration = self.read_unsigned_int()
99             if duration == 0:
100                 discontinuity_indicator = self.read_unsigned_char()
101             else:
102                 discontinuity_indicator = None
103             fragments.append({
104                 'first': first,
105                 'ts': first_ts,
106                 'duration': duration,
107                 'discontinuity_indicator': discontinuity_indicator,
108             })
109
110         return {
111             'fragments': fragments,
112         }
113
114     def read_abst(self):
115         # version
116         self.read_unsigned_char()
117         # flags
118         self.read(3)
119
120         self.read_unsigned_int()  # BootstrapinfoVersion
121         # Profile,Live,Update,Reserved
122         flags = self.read_unsigned_char()
123         live = flags & 0x20 != 0
124         # time scale
125         self.read_unsigned_int()
126         # CurrentMediaTime
127         self.read_unsigned_long_long()
128         # SmpteTimeCodeOffset
129         self.read_unsigned_long_long()
130
131         self.read_string()  # MovieIdentifier
132         server_count = self.read_unsigned_char()
133         # ServerEntryTable
134         for i in range(server_count):
135             self.read_string()
136         quality_count = self.read_unsigned_char()
137         # QualityEntryTable
138         for i in range(quality_count):
139             self.read_string()
140         # DrmData
141         self.read_string()
142         # MetaData
143         self.read_string()
144
145         segments_count = self.read_unsigned_char()
146         segments = []
147         for i in range(segments_count):
148             box_size, box_type, box_data = self.read_box_info()
149             assert box_type == b'asrt'
150             segment = FlvReader(box_data).read_asrt()
151             segments.append(segment)
152         fragments_run_count = self.read_unsigned_char()
153         fragments = []
154         for i in range(fragments_run_count):
155             box_size, box_type, box_data = self.read_box_info()
156             assert box_type == b'afrt'
157             fragments.append(FlvReader(box_data).read_afrt())
158
159         return {
160             'segments': segments,
161             'fragments': fragments,
162             'live': live,
163         }
164
165     def read_bootstrap_info(self):
166         total_size, box_type, box_data = self.read_box_info()
167         assert box_type == b'abst'
168         return FlvReader(box_data).read_abst()
169
170
171 def read_bootstrap_info(bootstrap_bytes):
172     return FlvReader(bootstrap_bytes).read_bootstrap_info()
173
174
175 def build_fragments_list(boot_info):
176     """ Return a list of (segment, fragment) for each fragment in the video """
177     res = []
178     segment_run_table = boot_info['segments'][0]
179     fragment_run_entry_table = boot_info['fragments'][0]['fragments']
180     first_frag_number = fragment_run_entry_table[0]['first']
181     fragments_counter = itertools.count(first_frag_number)
182     for segment, fragments_count in segment_run_table['segment_run']:
183         for _ in range(fragments_count):
184             res.append((segment, next(fragments_counter)))
185
186     if boot_info['live']:
187         res = res[-2:]
188
189     return res
190
191
192 def write_unsigned_int(stream, val):
193     stream.write(struct_pack('!I', val))
194
195
196 def write_unsigned_int_24(stream, val):
197     stream.write(struct_pack('!I', val)[1:])
198
199
200 def write_flv_header(stream):
201     """Writes the FLV header to stream"""
202     # FLV header
203     stream.write(b'FLV\x01')
204     stream.write(b'\x05')
205     stream.write(b'\x00\x00\x00\x09')
206     stream.write(b'\x00\x00\x00\x00')
207
208
209 def write_metadata_tag(stream, metadata):
210     """Writes optional metadata tag to stream"""
211     SCRIPT_TAG = b'\x12'
212     FLV_TAG_HEADER_LEN = 11
213
214     if metadata:
215         stream.write(SCRIPT_TAG)
216         write_unsigned_int_24(stream, len(metadata))
217         stream.write(b'\x00\x00\x00\x00\x00\x00\x00')
218         stream.write(metadata)
219         write_unsigned_int(stream, FLV_TAG_HEADER_LEN + len(metadata))
220
221
222 def _add_ns(prop):
223     return '{http://ns.adobe.com/f4m/1.0}%s' % prop
224
225
226 class F4mFD(FragmentFD):
227     """
228     A downloader for f4m manifests or AdobeHDS.
229     """
230
231     FD_NAME = 'f4m'
232
233     def _get_unencrypted_media(self, doc):
234         media = doc.findall(_add_ns('media'))
235         if not media:
236             self.report_error('No media found')
237         for e in (doc.findall(_add_ns('drmAdditionalHeader')) +
238                   doc.findall(_add_ns('drmAdditionalHeaderSet'))):
239             # If id attribute is missing it's valid for all media nodes
240             # without drmAdditionalHeaderId or drmAdditionalHeaderSetId attribute
241             if 'id' not in e.attrib:
242                 self.report_error('Missing ID in f4m DRM')
243         media = list(filter(lambda e: 'drmAdditionalHeaderId' not in e.attrib and
244                                       'drmAdditionalHeaderSetId' not in e.attrib,
245                             media))
246         if not media:
247             self.report_error('Unsupported DRM')
248         return media
249
250     def _get_bootstrap_from_url(self, bootstrap_url):
251         bootstrap = self.ydl.urlopen(bootstrap_url).read()
252         return read_bootstrap_info(bootstrap)
253
254     def _update_live_fragments(self, bootstrap_url, latest_fragment):
255         fragments_list = []
256         retries = 30
257         while (not fragments_list) and (retries > 0):
258             boot_info = self._get_bootstrap_from_url(bootstrap_url)
259             fragments_list = build_fragments_list(boot_info)
260             fragments_list = [f for f in fragments_list if f[1] > latest_fragment]
261             if not fragments_list:
262                 # Retry after a while
263                 time.sleep(5.0)
264                 retries -= 1
265
266         if not fragments_list:
267             self.report_error('Failed to update fragments')
268
269         return fragments_list
270
271     def _parse_bootstrap_node(self, node, base_url):
272         if node.text is None:
273             bootstrap_url = compat_urlparse.urljoin(
274                 base_url, node.attrib['url'])
275             boot_info = self._get_bootstrap_from_url(bootstrap_url)
276         else:
277             bootstrap_url = None
278             bootstrap = base64.b64decode(node.text.encode('ascii'))
279             boot_info = read_bootstrap_info(bootstrap)
280         return (boot_info, bootstrap_url)
281
282     def real_download(self, filename, info_dict):
283         man_url = info_dict['url']
284         requested_bitrate = info_dict.get('tbr')
285         self.to_screen('[%s] Downloading f4m manifest' % self.FD_NAME)
286         manifest = self.ydl.urlopen(man_url).read()
287
288         doc = etree.fromstring(manifest)
289         formats = [(int(f.attrib.get('bitrate', -1)), f)
290                    for f in self._get_unencrypted_media(doc)]
291         if requested_bitrate is None:
292             # get the best format
293             formats = sorted(formats, key=lambda f: f[0])
294             rate, media = formats[-1]
295         else:
296             rate, media = list(filter(
297                 lambda f: int(f[0]) == requested_bitrate, formats))[0]
298
299         base_url = compat_urlparse.urljoin(man_url, media.attrib['url'])
300         bootstrap_node = doc.find(_add_ns('bootstrapInfo'))
301         boot_info, bootstrap_url = self._parse_bootstrap_node(bootstrap_node, base_url)
302         live = boot_info['live']
303         metadata_node = media.find(_add_ns('metadata'))
304         if metadata_node is not None:
305             metadata = base64.b64decode(metadata_node.text.encode('ascii'))
306         else:
307             metadata = None
308
309         fragments_list = build_fragments_list(boot_info)
310         if self.params.get('test', False):
311             # We only download the first fragment
312             fragments_list = fragments_list[:1]
313         total_frags = len(fragments_list)
314         # For some akamai manifests we'll need to add a query to the fragment url
315         akamai_pv = xpath_text(doc, _add_ns('pv-2.0'))
316
317         ctx = {
318             'filename': filename,
319             'total_frags': total_frags,
320         }
321
322         self._prepare_frag_download(ctx)
323
324         dest_stream = ctx['dest_stream']
325
326         write_flv_header(dest_stream)
327         if not live:
328             write_metadata_tag(dest_stream, metadata)
329
330         self._start_frag_download(ctx)
331
332         frags_filenames = []
333         while fragments_list:
334             seg_i, frag_i = fragments_list.pop(0)
335             name = 'Seg%d-Frag%d' % (seg_i, frag_i)
336             url = base_url + name
337             if akamai_pv:
338                 url += '?' + akamai_pv.strip(';')
339             if info_dict.get('extra_param_to_segment_url'):
340                 url += info_dict.get('extra_param_to_segment_url')
341             frag_filename = '%s-%s' % (ctx['tmpfilename'], name)
342             try:
343                 success = ctx['dl'].download(frag_filename, {'url': url})
344                 if not success:
345                     return False
346                 (down, frag_sanitized) = sanitize_open(frag_filename, 'rb')
347                 down_data = down.read()
348                 down.close()
349                 reader = FlvReader(down_data)
350                 while True:
351                     _, box_type, box_data = reader.read_box_info()
352                     if box_type == b'mdat':
353                         dest_stream.write(box_data)
354                         break
355                 if live:
356                     os.remove(encodeFilename(frag_sanitized))
357                 else:
358                     frags_filenames.append(frag_sanitized)
359             except (compat_urllib_error.HTTPError, ) as err:
360                 if live and (err.code == 404 or err.code == 410):
361                     # We didn't keep up with the live window. Continue
362                     # with the next available fragment.
363                     msg = 'Fragment %d unavailable' % frag_i
364                     self.report_warning(msg)
365                     fragments_list = []
366                 else:
367                     raise
368
369             if not fragments_list and live and bootstrap_url:
370                 fragments_list = self._update_live_fragments(bootstrap_url, frag_i)
371                 total_frags += len(fragments_list)
372                 if fragments_list and (fragments_list[0][1] > frag_i + 1):
373                     msg = 'Missed %d fragments' % (fragments_list[0][1] - (frag_i + 1))
374                     self.report_warning(msg)
375
376         self._finish_frag_download(ctx)
377
378         for frag_file in frags_filenames:
379             os.remove(encodeFilename(frag_file))
380
381         return True