mirror of
https://github.com/google/bumble.git
synced 2026-04-16 00:25:31 +00:00
534 lines
20 KiB
Python
534 lines
20 KiB
Python
# Copyright 2023 Google LLC
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# https://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# Imports
|
|
# -----------------------------------------------------------------------------
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
|
|
from typing_extensions import Self
|
|
|
|
from bumble import core
|
|
|
|
|
|
# -----------------------------------------------------------------------------
|
|
class BitReader:
|
|
"""Simple but not optimized bit stream reader."""
|
|
|
|
data: bytes
|
|
bytes_position: int
|
|
bit_position: int
|
|
cache: int
|
|
bits_cached: int
|
|
|
|
def __init__(self, data: bytes):
|
|
self.data = data
|
|
self.byte_position = 0
|
|
self.bit_position = 0
|
|
self.cache = 0
|
|
self.bits_cached = 0
|
|
|
|
def read(self, bits: int) -> int:
|
|
""" "Read up to 32 bits."""
|
|
|
|
if bits > 32:
|
|
raise core.InvalidArgumentError('maximum read size is 32')
|
|
|
|
if self.bits_cached >= bits:
|
|
# We have enough bits.
|
|
self.bits_cached -= bits
|
|
self.bit_position += bits
|
|
return (self.cache >> self.bits_cached) & ((1 << bits) - 1)
|
|
|
|
# Read more cache, up to 32 bits
|
|
feed_bytes = self.data[self.byte_position : self.byte_position + 4]
|
|
feed_size = len(feed_bytes)
|
|
feed_int = int.from_bytes(feed_bytes, byteorder='big')
|
|
if 8 * feed_size + self.bits_cached < bits:
|
|
raise core.InvalidArgumentError('trying to read past the data')
|
|
self.byte_position += feed_size
|
|
|
|
# Combine the new cache and the old cache
|
|
cache = self.cache & ((1 << self.bits_cached) - 1)
|
|
new_bits = bits - self.bits_cached
|
|
self.bits_cached = 8 * feed_size - new_bits
|
|
result = (feed_int >> self.bits_cached) | (cache << new_bits)
|
|
self.cache = feed_int
|
|
|
|
self.bit_position += bits
|
|
return result
|
|
|
|
def read_bytes(self, count: int):
|
|
if self.bit_position + 8 * count > 8 * len(self.data):
|
|
raise core.InvalidArgumentError('not enough data')
|
|
|
|
if self.bit_position % 8:
|
|
# Not byte aligned
|
|
result = bytearray(count)
|
|
for i in range(count):
|
|
result[i] = self.read(8)
|
|
return bytes(result)
|
|
|
|
# Byte aligned
|
|
self.byte_position = self.bit_position // 8
|
|
self.bits_cached = 0
|
|
self.cache = 0
|
|
offset = self.bit_position // 8
|
|
self.bit_position += 8 * count
|
|
return self.data[offset : offset + count]
|
|
|
|
def bits_left(self) -> int:
|
|
return (8 * len(self.data)) - self.bit_position
|
|
|
|
def skip(self, bits: int) -> None:
|
|
# Slow, but simple...
|
|
while bits:
|
|
if bits > 32:
|
|
self.read(32)
|
|
bits -= 32
|
|
else:
|
|
self.read(bits)
|
|
break
|
|
|
|
|
|
# -----------------------------------------------------------------------------
|
|
class BitWriter:
|
|
"""Simple but not optimized bit stream writer."""
|
|
|
|
data: int
|
|
bit_count: int
|
|
|
|
def __init__(self) -> None:
|
|
self.data = 0
|
|
self.bit_count = 0
|
|
|
|
def write(self, value: int, bit_count: int) -> None:
|
|
self.data = (self.data << bit_count) | value
|
|
self.bit_count += bit_count
|
|
|
|
def write_bytes(self, data: bytes) -> None:
|
|
bit_count = 8 * len(data)
|
|
self.data = (self.data << bit_count) | int.from_bytes(data, 'big')
|
|
self.bit_count += bit_count
|
|
|
|
def __bytes__(self) -> bytes:
|
|
return (self.data << ((8 - (self.bit_count % 8)) % 8)).to_bytes(
|
|
(self.bit_count + 7) // 8, 'big'
|
|
)
|
|
|
|
|
|
# -----------------------------------------------------------------------------
|
|
class AacAudioRtpPacket:
|
|
"""AAC payload encapsulated in an RTP packet payload"""
|
|
|
|
audio_mux_element: AudioMuxElement
|
|
|
|
@staticmethod
|
|
def read_latm_value(reader: BitReader) -> int:
|
|
bytes_for_value = reader.read(2)
|
|
value = 0
|
|
for _ in range(bytes_for_value + 1):
|
|
value = value * 256 + reader.read(8)
|
|
return value
|
|
|
|
@staticmethod
|
|
def read_audio_object_type(reader: BitReader):
|
|
# GetAudioObjectType - ISO/EIC 14496-3 Table 1.16
|
|
audio_object_type = reader.read(5)
|
|
if audio_object_type == 31:
|
|
audio_object_type = 32 + reader.read(6)
|
|
|
|
return audio_object_type
|
|
|
|
@dataclass
|
|
class GASpecificConfig:
|
|
audio_object_type: int
|
|
# NOTE: other fields not supported
|
|
|
|
@classmethod
|
|
def from_bits(
|
|
cls, reader: BitReader, channel_configuration: int, audio_object_type: int
|
|
) -> Self:
|
|
# GASpecificConfig - ISO/EIC 14496-3 Table 4.1
|
|
reader.read(1) # frame_length_flag
|
|
depends_on_core_coder = reader.read(1)
|
|
if depends_on_core_coder:
|
|
reader.read(14) # core_coder_delay
|
|
extension_flag = reader.read(1)
|
|
if not channel_configuration:
|
|
raise core.InvalidPacketError('program_config_element not supported')
|
|
if audio_object_type in (6, 20):
|
|
reader.read(3) # layer_nr
|
|
if extension_flag:
|
|
if audio_object_type == 22:
|
|
reader.read(5) # num_of_sub_frame
|
|
reader.read(11) # layer_length
|
|
if audio_object_type in (17, 19, 20, 23):
|
|
reader.read(1) # aac_section_data_resilience_flags
|
|
reader.read(1) # aac_scale_factor_data_resilience_flags
|
|
reader.read(1) # aac_spectral_data_resilience_flags
|
|
extension_flag_3 = reader.read(1)
|
|
if extension_flag_3 == 1:
|
|
raise core.InvalidPacketError('extensionFlag3 == 1 not supported')
|
|
|
|
return cls(audio_object_type)
|
|
|
|
def to_bits(self, writer: BitWriter) -> None:
|
|
assert self.audio_object_type in (1, 2)
|
|
writer.write(0, 1) # frame_length_flag = 0
|
|
writer.write(0, 1) # depends_on_core_coder = 0
|
|
writer.write(0, 1) # extension_flag = 0
|
|
|
|
@dataclass
|
|
class AudioSpecificConfig:
|
|
audio_object_type: int
|
|
sampling_frequency_index: int
|
|
sampling_frequency: int
|
|
channel_configuration: int
|
|
ga_specific_config: AacAudioRtpPacket.GASpecificConfig
|
|
sbr_present_flag: int
|
|
ps_present_flag: int
|
|
extension_audio_object_type: int
|
|
extension_sampling_frequency_index: int
|
|
extension_sampling_frequency: int
|
|
extension_channel_configuration: int
|
|
|
|
SAMPLING_FREQUENCIES = [
|
|
96000,
|
|
88200,
|
|
64000,
|
|
48000,
|
|
44100,
|
|
32000,
|
|
24000,
|
|
22050,
|
|
16000,
|
|
12000,
|
|
11025,
|
|
8000,
|
|
7350,
|
|
]
|
|
|
|
@classmethod
|
|
def for_simple_aac(
|
|
cls,
|
|
audio_object_type: int,
|
|
sampling_frequency: int,
|
|
channel_configuration: int,
|
|
) -> Self:
|
|
if sampling_frequency not in cls.SAMPLING_FREQUENCIES:
|
|
raise ValueError(f'invalid sampling frequency {sampling_frequency}')
|
|
|
|
ga_specific_config = AacAudioRtpPacket.GASpecificConfig(audio_object_type)
|
|
|
|
return cls(
|
|
audio_object_type=audio_object_type,
|
|
sampling_frequency_index=cls.SAMPLING_FREQUENCIES.index(
|
|
sampling_frequency
|
|
),
|
|
sampling_frequency=sampling_frequency,
|
|
channel_configuration=channel_configuration,
|
|
ga_specific_config=ga_specific_config,
|
|
sbr_present_flag=0,
|
|
ps_present_flag=0,
|
|
extension_audio_object_type=0,
|
|
extension_sampling_frequency_index=0,
|
|
extension_sampling_frequency=0,
|
|
extension_channel_configuration=0,
|
|
)
|
|
|
|
@classmethod
|
|
def from_bits(cls, reader: BitReader) -> Self:
|
|
# AudioSpecificConfig - ISO/EIC 14496-3 Table 1.15
|
|
audio_object_type = AacAudioRtpPacket.read_audio_object_type(reader)
|
|
sampling_frequency_index = reader.read(4)
|
|
if sampling_frequency_index == 0xF:
|
|
sampling_frequency = reader.read(24)
|
|
else:
|
|
sampling_frequency = cls.SAMPLING_FREQUENCIES[sampling_frequency_index]
|
|
channel_configuration = reader.read(4)
|
|
sbr_present_flag = 0
|
|
ps_present_flag = 0
|
|
extension_sampling_frequency_index = 0
|
|
extension_sampling_frequency = 0
|
|
extension_channel_configuration = 0
|
|
extension_audio_object_type = 0
|
|
if audio_object_type in (5, 29):
|
|
extension_audio_object_type = 5
|
|
sbr_present_flag = 1
|
|
if audio_object_type == 29:
|
|
ps_present_flag = 1
|
|
extension_sampling_frequency_index = reader.read(4)
|
|
if extension_sampling_frequency_index == 0xF:
|
|
extension_sampling_frequency = reader.read(24)
|
|
else:
|
|
extension_sampling_frequency = cls.SAMPLING_FREQUENCIES[
|
|
extension_sampling_frequency_index
|
|
]
|
|
audio_object_type = AacAudioRtpPacket.read_audio_object_type(reader)
|
|
if audio_object_type == 22:
|
|
extension_channel_configuration = reader.read(4)
|
|
|
|
if audio_object_type in (1, 2, 3, 4, 6, 7, 17, 19, 20, 21, 22, 23):
|
|
ga_specific_config = AacAudioRtpPacket.GASpecificConfig.from_bits(
|
|
reader, channel_configuration, audio_object_type
|
|
)
|
|
else:
|
|
raise core.InvalidPacketError(
|
|
f'audioObjectType {audio_object_type} not supported'
|
|
)
|
|
|
|
# if self.extension_audio_object_type != 5 and bits_to_decode >= 16:
|
|
# sync_extension_type = reader.read(11)
|
|
# if sync_extension_type == 0x2B7:
|
|
# self.extension_audio_object_type = AacAudioRtpPacket.audio_object_type(reader)
|
|
# if self.extension_audio_object_type == 5:
|
|
# self.sbr_present_flag = reader.read(1)
|
|
# if self.sbr_present_flag:
|
|
# self.extension_sampling_frequency_index = reader.read(4)
|
|
# if self.extension_sampling_frequency_index == 0xF:
|
|
# self.extension_sampling_frequency = reader.read(24)
|
|
# else:
|
|
# self.extension_sampling_frequency = self.SAMPLING_FREQUENCIES[self.extension_sampling_frequency_index]
|
|
# if bits_to_decode >= 12:
|
|
# sync_extension_type = reader.read(11)
|
|
# if sync_extension_type == 0x548:
|
|
# self.ps_present_flag = reader.read(1)
|
|
# elif self.extension_audio_object_type == 22:
|
|
# self.sbr_present_flag = reader.read(1)
|
|
# if self.sbr_present_flag:
|
|
# self.extension_sampling_frequency_index = reader.read(4)
|
|
# if self.extension_sampling_frequency_index == 0xF:
|
|
# self.extension_sampling_frequency = reader.read(24)
|
|
# else:
|
|
# self.extension_sampling_frequency = self.SAMPLING_FREQUENCIES[self.extension_sampling_frequency_index]
|
|
# self.extension_channel_configuration = reader.read(4)
|
|
|
|
return cls(
|
|
audio_object_type,
|
|
sampling_frequency_index,
|
|
sampling_frequency,
|
|
channel_configuration,
|
|
ga_specific_config,
|
|
sbr_present_flag,
|
|
ps_present_flag,
|
|
extension_audio_object_type,
|
|
extension_sampling_frequency_index,
|
|
extension_sampling_frequency,
|
|
extension_channel_configuration,
|
|
)
|
|
|
|
def to_bits(self, writer: BitWriter) -> None:
|
|
if self.sampling_frequency_index >= 15:
|
|
raise ValueError(
|
|
f"unsupported sampling frequency index {self.sampling_frequency_index}"
|
|
)
|
|
|
|
if self.audio_object_type not in (1, 2):
|
|
raise ValueError(
|
|
f"unsupported audio object type {self.audio_object_type} "
|
|
)
|
|
|
|
writer.write(self.audio_object_type, 5)
|
|
writer.write(self.sampling_frequency_index, 4)
|
|
writer.write(self.channel_configuration, 4)
|
|
self.ga_specific_config.to_bits(writer)
|
|
|
|
@dataclass
|
|
class StreamMuxConfig:
|
|
other_data_present: int
|
|
other_data_len_bits: int
|
|
audio_specific_config: AacAudioRtpPacket.AudioSpecificConfig
|
|
|
|
@classmethod
|
|
def from_bits(cls, reader: BitReader) -> Self:
|
|
# StreamMuxConfig - ISO/EIC 14496-3 Table 1.42
|
|
audio_mux_version = reader.read(1)
|
|
if audio_mux_version == 1:
|
|
audio_mux_version_a = reader.read(1)
|
|
else:
|
|
audio_mux_version_a = 0
|
|
if audio_mux_version_a != 0:
|
|
raise core.InvalidPacketError('audioMuxVersionA != 0 not supported')
|
|
if audio_mux_version == 1:
|
|
AacAudioRtpPacket.read_latm_value(reader) # tara_buffer_fullness
|
|
# stream_cnt = 0
|
|
reader.read(1) # all_streams_same_time_framing
|
|
reader.read(6) # num_sub_frames
|
|
num_program = reader.read(4)
|
|
if num_program != 0:
|
|
raise core.InvalidPacketError('num_program != 0 not supported')
|
|
num_layer = reader.read(3)
|
|
if num_layer != 0:
|
|
raise core.InvalidPacketError('num_layer != 0 not supported')
|
|
if audio_mux_version == 0:
|
|
audio_specific_config = AacAudioRtpPacket.AudioSpecificConfig.from_bits(
|
|
reader
|
|
)
|
|
else:
|
|
asc_len = AacAudioRtpPacket.read_latm_value(reader)
|
|
marker = reader.bit_position
|
|
audio_specific_config = AacAudioRtpPacket.AudioSpecificConfig.from_bits(
|
|
reader
|
|
)
|
|
audio_specific_config_len = reader.bit_position - marker
|
|
if asc_len < audio_specific_config_len:
|
|
raise core.InvalidPacketError('audio_specific_config_len > asc_len')
|
|
asc_len -= audio_specific_config_len
|
|
reader.skip(asc_len)
|
|
frame_length_type = reader.read(3)
|
|
if frame_length_type == 0:
|
|
reader.read(8) # latm_buffer_fullness
|
|
elif frame_length_type == 1:
|
|
reader.read(9) # frame_length
|
|
else:
|
|
raise core.InvalidPacketError(
|
|
f'frame_length_type {frame_length_type} not supported'
|
|
)
|
|
|
|
other_data_present = reader.read(1)
|
|
other_data_len_bits = 0
|
|
if other_data_present:
|
|
if audio_mux_version == 1:
|
|
other_data_len_bits = AacAudioRtpPacket.read_latm_value(reader)
|
|
else:
|
|
while True:
|
|
other_data_len_bits *= 256
|
|
other_data_len_esc = reader.read(1)
|
|
other_data_len_bits += reader.read(8)
|
|
if other_data_len_esc == 0:
|
|
break
|
|
crc_check_present = reader.read(1)
|
|
if crc_check_present:
|
|
reader.read(8) # crc_checksum
|
|
|
|
return cls(other_data_present, other_data_len_bits, audio_specific_config)
|
|
|
|
def to_bits(self, writer: BitWriter) -> None:
|
|
writer.write(0, 1) # audioMuxVersion = 0
|
|
writer.write(1, 1) # allStreamsSameTimeFraming = 1
|
|
writer.write(0, 6) # numSubFrames = 0
|
|
writer.write(0, 4) # numProgram = 0
|
|
writer.write(0, 3) # numLayer = 0
|
|
self.audio_specific_config.to_bits(writer)
|
|
writer.write(0, 3) # frameLengthType = 0
|
|
writer.write(0, 8) # latmBufferFullness = 0
|
|
writer.write(0, 1) # otherDataPresent = 0
|
|
writer.write(0, 1) # crcCheckPresent = 0
|
|
|
|
@dataclass
|
|
class AudioMuxElement:
|
|
stream_mux_config: AacAudioRtpPacket.StreamMuxConfig
|
|
payload: bytes
|
|
|
|
@classmethod
|
|
def from_bits(cls, reader: BitReader) -> Self:
|
|
# AudioMuxElement - ISO/EIC 14496-3 Table 1.41
|
|
# (only supports mux_config_present=1)
|
|
use_same_stream_mux = reader.read(1)
|
|
if use_same_stream_mux:
|
|
raise core.InvalidPacketError('useSameStreamMux == 1 not supported')
|
|
stream_mux_config = AacAudioRtpPacket.StreamMuxConfig.from_bits(reader)
|
|
|
|
# We only support:
|
|
# allStreamsSameTimeFraming == 1
|
|
# audioMuxVersionA == 0,
|
|
# numProgram == 0
|
|
# numSubFrames == 0
|
|
# numLayer == 0
|
|
|
|
mux_slot_length_bytes = 0
|
|
while True:
|
|
tmp = reader.read(8)
|
|
mux_slot_length_bytes += tmp
|
|
if tmp != 255:
|
|
break
|
|
|
|
payload = reader.read_bytes(mux_slot_length_bytes)
|
|
|
|
if stream_mux_config.other_data_present:
|
|
reader.skip(stream_mux_config.other_data_len_bits)
|
|
|
|
# ByteAlign
|
|
while reader.bit_position % 8:
|
|
reader.read(1)
|
|
|
|
return cls(stream_mux_config, payload)
|
|
|
|
def to_bits(self, writer: BitWriter) -> None:
|
|
writer.write(0, 1) # useSameStreamMux = 0
|
|
self.stream_mux_config.to_bits(writer)
|
|
mux_slot_length_bytes = len(self.payload)
|
|
while mux_slot_length_bytes > 255:
|
|
writer.write(255, 8)
|
|
mux_slot_length_bytes -= 255
|
|
writer.write(mux_slot_length_bytes, 8)
|
|
if mux_slot_length_bytes == 255:
|
|
writer.write(0, 8)
|
|
writer.write_bytes(self.payload)
|
|
|
|
@classmethod
|
|
def from_bytes(cls, data: bytes) -> Self:
|
|
# Parse the bit stream
|
|
reader = BitReader(data)
|
|
return cls(cls.AudioMuxElement.from_bits(reader))
|
|
|
|
@classmethod
|
|
def for_simple_aac(
|
|
cls, sampling_frequency: int, channel_configuration: int, payload: bytes
|
|
) -> Self:
|
|
audio_specific_config = cls.AudioSpecificConfig.for_simple_aac(
|
|
2, sampling_frequency, channel_configuration
|
|
)
|
|
stream_mux_config = cls.StreamMuxConfig(0, 0, audio_specific_config)
|
|
audio_mux_element = cls.AudioMuxElement(stream_mux_config, payload)
|
|
|
|
return cls(audio_mux_element)
|
|
|
|
def to_adts(self):
|
|
# pylint: disable=line-too-long
|
|
sampling_frequency_index = (
|
|
self.audio_mux_element.stream_mux_config.audio_specific_config.sampling_frequency_index
|
|
)
|
|
channel_configuration = (
|
|
self.audio_mux_element.stream_mux_config.audio_specific_config.channel_configuration
|
|
)
|
|
frame_size = len(self.audio_mux_element.payload)
|
|
return (
|
|
bytes(
|
|
[
|
|
0xFF,
|
|
0xF1, # 0xF9 (MPEG2)
|
|
0x40
|
|
| (sampling_frequency_index << 2)
|
|
| (channel_configuration >> 2),
|
|
((channel_configuration & 0x3) << 6) | ((frame_size + 7) >> 11),
|
|
((frame_size + 7) >> 3) & 0xFF,
|
|
(((frame_size + 7) << 5) & 0xFF) | 0x1F,
|
|
0xFC,
|
|
]
|
|
)
|
|
+ self.audio_mux_element.payload
|
|
)
|
|
|
|
def __init__(self, audio_mux_element: AudioMuxElement) -> None:
|
|
self.audio_mux_element = audio_mux_element
|
|
|
|
def __bytes__(self) -> bytes:
|
|
writer = BitWriter()
|
|
self.audio_mux_element.to_bits(writer)
|
|
return bytes(writer)
|