From e5ed6236e350cfb8539c04eb7b041b761b7b6314 Mon Sep 17 00:00:00 2001
From: rlaphoenix <rlaphoenix@pm.me>
Date: Sun, 26 Mar 2023 09:31:19 +0100
Subject: [PATCH] Re-implement Segment.key, having the last key within keys

This keeps compatibility and fixes a lot of the existing tests. However, new tests should likely be implemented for testing HLS playlists with multiple preceeding keys.
---
 m3u8/model.py  | 9 +++++++--
 m3u8/parser.py | 6 ++++++
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/m3u8/model.py b/m3u8/model.py
index 84e8f1e8..4423ff44 100644
--- a/m3u8/model.py
+++ b/m3u8/model.py
@@ -172,6 +172,7 @@ def _initialize_attributes(self):
                 keyobjects=[
                     find_key(segment_key, self.keys)
                     for segment_key in segment['keys']],
+                keyobject=find_key(segment.get('key', {}), self.keys),
                 **segment)
             for segment in self.data.get('segments', [])
         ])
@@ -436,7 +437,10 @@ class Segment(BasePathMixin):
       byterange attribute from EXT-X-BYTERANGE parameter
 
     `keys`
-      Keys used to encrypt the segment (EXT-X-KEY)
+      Keys used to encrypt the segment (list of EXT-X-KEY)
+
+    `key`
+      Last Key within keys used to encrypt the segment (EXT-X-KEY)
 
     `parts`
       partial segments that make up this segment
@@ -454,7 +458,7 @@ class Segment(BasePathMixin):
     def __init__(self, uri=None, base_uri=None, program_date_time=None, current_program_date_time=None,
                  duration=None, title=None, bitrate=None, byterange=None, cue_out=False,
                  cue_out_start=False, cue_in=False, discontinuity=False, keys=None, scte35=None,
-                 oatcls_scte35=None, scte35_duration=None, scte35_elapsedtime=None, asset_metadata=None,
+                 oatcls_scte35=None, scte35_duration=None, scte35_elapsedtime=None, asset_metadata=None, keyobject=None,
                  keyobjects=None, parts=None, init_section=None, dateranges=None, gap_tag=None,
                  media_sequence=None, custom_parser_values=None):
         self.media_sequence = media_sequence
@@ -476,6 +480,7 @@ def __init__(self, uri=None, base_uri=None, program_date_time=None, current_prog
         self.scte35_elapsedtime = scte35_elapsedtime
         self.asset_metadata = asset_metadata
         self.keys = keyobjects or []
+        self.key = keyobject
         self.parts = PartialSegmentList( [ PartialSegment(base_uri=self._base_uri, **partial) for partial in parts ] if parts else [] )
         if init_section is not None:
             self.init_section = InitializationSection(self._base_uri, **init_section)
diff --git a/m3u8/parser.py b/m3u8/parser.py
index 47e615db..76e2df77 100644
--- a/m3u8/parser.py
+++ b/m3u8/parser.py
@@ -64,6 +64,7 @@ def parse(content, strict=False, custom_tags_parser=None):
         'expect_segment': False,
         'expect_playlist': False,
         'current_keys': [],
+        'current_key': None,
         'current_segment_map': None,
     }
 
@@ -97,6 +98,7 @@ def parse(content, strict=False, custom_tags_parser=None):
         elif line.startswith(protocol.ext_x_discontinuity_sequence):
             _parse_simple_parameter(line, data, int)
             state['current_keys'].clear()
+            state['current_key'] = None
 
         elif line.startswith(protocol.ext_x_program_date_time):
             _, program_date_time = _parse_simple_parameter_raw_value(line, cast_date_time)
@@ -138,6 +140,7 @@ def parse(content, strict=False, custom_tags_parser=None):
         elif line.startswith(protocol.ext_x_key):
             key = _parse_key(line)
             state['current_keys'].append(key)
+            state['current_key'] = key
             if key not in data['keys']:
                 data['keys'].append(key)
 
@@ -225,6 +228,7 @@ def parse(content, strict=False, custom_tags_parser=None):
             _parse_ts_chunk(line, data, state)
             state['expect_segment'] = False
             state['current_keys'].clear()
+            state['current_key'] = None
 
         elif state['expect_playlist']:
             _parse_variant_playlist(line, data, state)
@@ -284,6 +288,8 @@ def _parse_ts_chunk(line, data, state):
     segment['asset_metadata'] = scte_op('asset_metadata', None)
     segment['discontinuity'] = state.pop('discontinuity', False)
     segment['keys'] = copy(state['current_keys'])
+    if state.get('current_key'):
+        segment['key'] = state['current_key']
     if not state['current_keys']:
         # For unencrypted segments, the initial key would be None
         if None not in data['keys']: