]> jfr.im git - yt-dlp.git/blob - yt_dlp/downloader/ism.py
[cleanup] Add more ruff rules (#10149)
[yt-dlp.git] / yt_dlp / downloader / ism.py
1 import binascii
2 import io
3 import struct
4 import time
5
6 from .fragment import FragmentFD
7 from ..networking.exceptions import HTTPError
8 from ..utils import RetryManager
9
10 u8 = struct.Struct('>B')
11 u88 = struct.Struct('>Bx')
12 u16 = struct.Struct('>H')
13 u1616 = struct.Struct('>Hxx')
14 u32 = struct.Struct('>I')
15 u64 = struct.Struct('>Q')
16
17 s88 = struct.Struct('>bx')
18 s16 = struct.Struct('>h')
19 s1616 = struct.Struct('>hxx')
20 s32 = struct.Struct('>i')
21
22 unity_matrix = (s32.pack(0x10000) + s32.pack(0) * 3) * 2 + s32.pack(0x40000000)
23
24 TRACK_ENABLED = 0x1
25 TRACK_IN_MOVIE = 0x2
26 TRACK_IN_PREVIEW = 0x4
27
28 SELF_CONTAINED = 0x1
29
30
31 def box(box_type, payload):
32 return u32.pack(8 + len(payload)) + box_type + payload
33
34
35 def full_box(box_type, version, flags, payload):
36 return box(box_type, u8.pack(version) + u32.pack(flags)[1:] + payload)
37
38
39 def write_piff_header(stream, params):
40 track_id = params['track_id']
41 fourcc = params['fourcc']
42 duration = params['duration']
43 timescale = params.get('timescale', 10000000)
44 language = params.get('language', 'und')
45 height = params.get('height', 0)
46 width = params.get('width', 0)
47 stream_type = params['stream_type']
48 creation_time = modification_time = int(time.time())
49
50 ftyp_payload = b'isml' # major brand
51 ftyp_payload += u32.pack(1) # minor version
52 ftyp_payload += b'piff' + b'iso2' # compatible brands
53 stream.write(box(b'ftyp', ftyp_payload)) # File Type Box
54
55 mvhd_payload = u64.pack(creation_time)
56 mvhd_payload += u64.pack(modification_time)
57 mvhd_payload += u32.pack(timescale)
58 mvhd_payload += u64.pack(duration)
59 mvhd_payload += s1616.pack(1) # rate
60 mvhd_payload += s88.pack(1) # volume
61 mvhd_payload += u16.pack(0) # reserved
62 mvhd_payload += u32.pack(0) * 2 # reserved
63 mvhd_payload += unity_matrix
64 mvhd_payload += u32.pack(0) * 6 # pre defined
65 mvhd_payload += u32.pack(0xffffffff) # next track id
66 moov_payload = full_box(b'mvhd', 1, 0, mvhd_payload) # Movie Header Box
67
68 tkhd_payload = u64.pack(creation_time)
69 tkhd_payload += u64.pack(modification_time)
70 tkhd_payload += u32.pack(track_id) # track id
71 tkhd_payload += u32.pack(0) # reserved
72 tkhd_payload += u64.pack(duration)
73 tkhd_payload += u32.pack(0) * 2 # reserved
74 tkhd_payload += s16.pack(0) # layer
75 tkhd_payload += s16.pack(0) # alternate group
76 tkhd_payload += s88.pack(1 if stream_type == 'audio' else 0) # volume
77 tkhd_payload += u16.pack(0) # reserved
78 tkhd_payload += unity_matrix
79 tkhd_payload += u1616.pack(width)
80 tkhd_payload += u1616.pack(height)
81 trak_payload = full_box(b'tkhd', 1, TRACK_ENABLED | TRACK_IN_MOVIE | TRACK_IN_PREVIEW, tkhd_payload) # Track Header Box
82
83 mdhd_payload = u64.pack(creation_time)
84 mdhd_payload += u64.pack(modification_time)
85 mdhd_payload += u32.pack(timescale)
86 mdhd_payload += u64.pack(duration)
87 mdhd_payload += u16.pack(((ord(language[0]) - 0x60) << 10) | ((ord(language[1]) - 0x60) << 5) | (ord(language[2]) - 0x60))
88 mdhd_payload += u16.pack(0) # pre defined
89 mdia_payload = full_box(b'mdhd', 1, 0, mdhd_payload) # Media Header Box
90
91 hdlr_payload = u32.pack(0) # pre defined
92 if stream_type == 'audio': # handler type
93 hdlr_payload += b'soun'
94 hdlr_payload += u32.pack(0) * 3 # reserved
95 hdlr_payload += b'SoundHandler\0' # name
96 elif stream_type == 'video':
97 hdlr_payload += b'vide'
98 hdlr_payload += u32.pack(0) * 3 # reserved
99 hdlr_payload += b'VideoHandler\0' # name
100 elif stream_type == 'text':
101 hdlr_payload += b'subt'
102 hdlr_payload += u32.pack(0) * 3 # reserved
103 hdlr_payload += b'SubtitleHandler\0' # name
104 else:
105 assert False
106 mdia_payload += full_box(b'hdlr', 0, 0, hdlr_payload) # Handler Reference Box
107
108 if stream_type == 'audio':
109 smhd_payload = s88.pack(0) # balance
110 smhd_payload += u16.pack(0) # reserved
111 media_header_box = full_box(b'smhd', 0, 0, smhd_payload) # Sound Media Header
112 elif stream_type == 'video':
113 vmhd_payload = u16.pack(0) # graphics mode
114 vmhd_payload += u16.pack(0) * 3 # opcolor
115 media_header_box = full_box(b'vmhd', 0, 1, vmhd_payload) # Video Media Header
116 elif stream_type == 'text':
117 media_header_box = full_box(b'sthd', 0, 0, b'') # Subtitle Media Header
118 else:
119 assert False
120 minf_payload = media_header_box
121
122 dref_payload = u32.pack(1) # entry count
123 dref_payload += full_box(b'url ', 0, SELF_CONTAINED, b'') # Data Entry URL Box
124 dinf_payload = full_box(b'dref', 0, 0, dref_payload) # Data Reference Box
125 minf_payload += box(b'dinf', dinf_payload) # Data Information Box
126
127 stsd_payload = u32.pack(1) # entry count
128
129 sample_entry_payload = u8.pack(0) * 6 # reserved
130 sample_entry_payload += u16.pack(1) # data reference index
131 if stream_type == 'audio':
132 sample_entry_payload += u32.pack(0) * 2 # reserved
133 sample_entry_payload += u16.pack(params.get('channels', 2))
134 sample_entry_payload += u16.pack(params.get('bits_per_sample', 16))
135 sample_entry_payload += u16.pack(0) # pre defined
136 sample_entry_payload += u16.pack(0) # reserved
137 sample_entry_payload += u1616.pack(params['sampling_rate'])
138
139 if fourcc == 'AACL':
140 sample_entry_box = box(b'mp4a', sample_entry_payload)
141 if fourcc == 'EC-3':
142 sample_entry_box = box(b'ec-3', sample_entry_payload)
143 elif stream_type == 'video':
144 sample_entry_payload += u16.pack(0) # pre defined
145 sample_entry_payload += u16.pack(0) # reserved
146 sample_entry_payload += u32.pack(0) * 3 # pre defined
147 sample_entry_payload += u16.pack(width)
148 sample_entry_payload += u16.pack(height)
149 sample_entry_payload += u1616.pack(0x48) # horiz resolution 72 dpi
150 sample_entry_payload += u1616.pack(0x48) # vert resolution 72 dpi
151 sample_entry_payload += u32.pack(0) # reserved
152 sample_entry_payload += u16.pack(1) # frame count
153 sample_entry_payload += u8.pack(0) * 32 # compressor name
154 sample_entry_payload += u16.pack(0x18) # depth
155 sample_entry_payload += s16.pack(-1) # pre defined
156
157 codec_private_data = binascii.unhexlify(params['codec_private_data'].encode())
158 if fourcc in ('H264', 'AVC1'):
159 sps, pps = codec_private_data.split(u32.pack(1))[1:]
160 avcc_payload = u8.pack(1) # configuration version
161 avcc_payload += sps[1:4] # avc profile indication + profile compatibility + avc level indication
162 avcc_payload += u8.pack(0xfc | (params.get('nal_unit_length_field', 4) - 1)) # complete representation (1) + reserved (11111) + length size minus one
163 avcc_payload += u8.pack(1) # reserved (0) + number of sps (0000001)
164 avcc_payload += u16.pack(len(sps))
165 avcc_payload += sps
166 avcc_payload += u8.pack(1) # number of pps
167 avcc_payload += u16.pack(len(pps))
168 avcc_payload += pps
169 sample_entry_payload += box(b'avcC', avcc_payload) # AVC Decoder Configuration Record
170 sample_entry_box = box(b'avc1', sample_entry_payload) # AVC Simple Entry
171 else:
172 assert False
173 elif stream_type == 'text':
174 if fourcc == 'TTML':
175 sample_entry_payload += b'http://www.w3.org/ns/ttml\0' # namespace
176 sample_entry_payload += b'\0' # schema location
177 sample_entry_payload += b'\0' # auxilary mime types(??)
178 sample_entry_box = box(b'stpp', sample_entry_payload)
179 else:
180 assert False
181 else:
182 assert False
183 stsd_payload += sample_entry_box
184
185 stbl_payload = full_box(b'stsd', 0, 0, stsd_payload) # Sample Description Box
186
187 stts_payload = u32.pack(0) # entry count
188 stbl_payload += full_box(b'stts', 0, 0, stts_payload) # Decoding Time to Sample Box
189
190 stsc_payload = u32.pack(0) # entry count
191 stbl_payload += full_box(b'stsc', 0, 0, stsc_payload) # Sample To Chunk Box
192
193 stco_payload = u32.pack(0) # entry count
194 stbl_payload += full_box(b'stco', 0, 0, stco_payload) # Chunk Offset Box
195
196 minf_payload += box(b'stbl', stbl_payload) # Sample Table Box
197
198 mdia_payload += box(b'minf', minf_payload) # Media Information Box
199
200 trak_payload += box(b'mdia', mdia_payload) # Media Box
201
202 moov_payload += box(b'trak', trak_payload) # Track Box
203
204 mehd_payload = u64.pack(duration)
205 mvex_payload = full_box(b'mehd', 1, 0, mehd_payload) # Movie Extends Header Box
206
207 trex_payload = u32.pack(track_id) # track id
208 trex_payload += u32.pack(1) # default sample description index
209 trex_payload += u32.pack(0) # default sample duration
210 trex_payload += u32.pack(0) # default sample size
211 trex_payload += u32.pack(0) # default sample flags
212 mvex_payload += full_box(b'trex', 0, 0, trex_payload) # Track Extends Box
213
214 moov_payload += box(b'mvex', mvex_payload) # Movie Extends Box
215 stream.write(box(b'moov', moov_payload)) # Movie Box
216
217
218 def extract_box_data(data, box_sequence):
219 data_reader = io.BytesIO(data)
220 while True:
221 box_size = u32.unpack(data_reader.read(4))[0]
222 box_type = data_reader.read(4)
223 if box_type == box_sequence[0]:
224 box_data = data_reader.read(box_size - 8)
225 if len(box_sequence) == 1:
226 return box_data
227 return extract_box_data(box_data, box_sequence[1:])
228 data_reader.seek(box_size - 8, 1)
229
230
231 class IsmFD(FragmentFD):
232 """
233 Download segments in a ISM manifest
234 """
235
236 def real_download(self, filename, info_dict):
237 segments = info_dict['fragments'][:1] if self.params.get(
238 'test', False) else info_dict['fragments']
239
240 ctx = {
241 'filename': filename,
242 'total_frags': len(segments),
243 }
244
245 self._prepare_and_start_frag_download(ctx, info_dict)
246
247 extra_state = ctx.setdefault('extra_state', {
248 'ism_track_written': False,
249 })
250
251 skip_unavailable_fragments = self.params.get('skip_unavailable_fragments', True)
252
253 frag_index = 0
254 for segment in segments:
255 frag_index += 1
256 if frag_index <= ctx['fragment_index']:
257 continue
258
259 retry_manager = RetryManager(self.params.get('fragment_retries'), self.report_retry,
260 frag_index=frag_index, fatal=not skip_unavailable_fragments)
261 for retry in retry_manager:
262 try:
263 success = self._download_fragment(ctx, segment['url'], info_dict)
264 if not success:
265 return False
266 frag_content = self._read_fragment(ctx)
267
268 if not extra_state['ism_track_written']:
269 tfhd_data = extract_box_data(frag_content, [b'moof', b'traf', b'tfhd'])
270 info_dict['_download_params']['track_id'] = u32.unpack(tfhd_data[4:8])[0]
271 write_piff_header(ctx['dest_stream'], info_dict['_download_params'])
272 extra_state['ism_track_written'] = True
273 self._append_fragment(ctx, frag_content)
274 except HTTPError as err:
275 retry.error = err
276 continue
277
278 if retry_manager.error:
279 if not skip_unavailable_fragments:
280 return False
281 self.report_skip_fragment(frag_index)
282
283 return self._finish_frag_download(ctx, info_dict)