Repository: streamlink/streamlink
Pull request #2821 (2020-03-10 18:08:13)

add templating for hls-segment-key-uri option
This will replace the host in the segment URI with the supplied HOST

tarkah: satisfies #2806 
codecov[bot]: # [Codecov](https://codecov.io/gh/streamlink/streamlink/pull/2821?src=pr&el=h1) Report
> Merging [#2821](https://codecov.io/gh/streamlink/streamlink/pull/2821?src=pr&el=desc) into [master](https://codecov.io/gh/streamlink/streamlink/commit/8ef3c73480f172f1535d549a1a6c85c7da1c0c1f?src=pr&el=desc) will **decrease** coverage by `1.55%`.
> The diff coverage is `70%`.


```diff
@@            Coverage Diff             @@
##           master    #2821      +/-   ##
==========================================
- Coverage   52.93%   51.38%   -1.56%     
==========================================
  Files         248      248              
  Lines       15537    15195     -342     
==========================================
- Hits         8225     7808     -417     
- Misses       7312     7387      +75
```

bastimeyer: You should import urlparse from `streamlink.compat`:
https://github.com/streamlink/streamlink/blob/8ef3c73480f172f1535d549a1a6c85c7da1c0c1f/src/streamlink/compat.py#L35-L43

Regarding the implementation, why do we need two parameters for the replacement of segment keys? I think the already existing one could be improved, so that it can accept a string with `{variables}` instead, which would allow you to replace only the hostname, if that's what you need. Otherwise we could add parameters for each part of a URL, which would be nonsense.

Tests also have to be added before changes like this can be merged.
tarkah: Thanks for the feedback! I saw the 2.7 test fail for that same reason and it has been updated in latest push.

The new option doesn't specify what the _exact_ URI will be, since the segment URI can change throughout a stream. We won't know ahead of time what the full `/path?query` in the URI is since it will change. So not exactly sure how that can be handled in a templated fashion, unless I'm not understanding your suggestion.

I'll definitely add a test for this once we land on the optimal solution.
back-to: use only one command, **don't add a 2nd**

---

here is some unfinished code, might help you with it

```py
from base64 import b64encode

from streamlink.compat import urlparse

from streamlink.utils import LazyFormatter


        if self.key_uri_override:
            p = urlparse(key.uri)
            key_uri = LazyFormatter.format(
                self.key_uri_override,
                url=key.uri,
                url_base64=b64encode(key.uri.encode('utf-8')).decode('utf-8'),
                scheme=p.scheme,
                netloc=p.netloc,
                path=p.path,
                query=p.query,
            )
        else:
            key_uri = key.uri
```

this would allow `{url}{url_base64}{scheme}{netloc}{path}{query}` to be reused
tarkah: Ahh, very nifty! I'll leverage LazyFormatter, thanks!
tarkah: @back-to That worked perfectly. I still need to incorporate a test for this.
tarkah: How would you guys write a test for this? The HLS Stream is already well tested, this change just impacts the key URI that is requested, and an error is logged when that decryption fails. Should the functionality of `LazyFormatter` just be tested?

I've tested `--hls-segment-key-uri` with various templates at runtime and they all request the appropriate URI.
bastimeyer: You could add a test like this. Not sure if it's a bit too much though. Also not sure about the `aesKeyInvalid` line (getting this to work in Python 2 and 3 at the same time is weird), otherwise I'd add a commit to this PR myself.

```diff
diff --git a/tests/streams/test_hls.py b/tests/streams/test_hls.py
index 61db631e..26b35733 100644
--- a/tests/streams/test_hls.py
+++ b/tests/streams/test_hls.py
@@ -59,13 +59,14 @@ class TestHLS(unittest.TestCase):
 
         return playlist + playlistEnd
 
-    def start_streamlink(self, masterPlaylist, kwargs=None):
+    def start_streamlink(self, masterPlaylist, hls_segment_key_uri=None, kwargs=None):
         kwargs = kwargs or {}
         log.info("Executing streamlink")
         streamlink = Streamlink()
 
         # Set to default value to avoid a test fail if the default change
         streamlink.set_option("hls-live-edge", 3)
+        streamlink.set_option("hls-segment-key-uri", hls_segment_key_uri)
 
         masterStream = hls.HLSStream.parse_variant_playlist(streamlink, masterPlaylist, **kwargs)
         stream = masterStream["1080p (source)"].open()
@@ -86,13 +87,13 @@ class TestHLS(unittest.TestCase):
 
             # Start streamlink on the generated stream
             streamlinkResult = self.start_streamlink("http://mocked/path/master.m3u8",
-                                                     {'start_offset': 1, 'duration': 1})
+                                                     kwargs={'start_offset': 1, 'duration': 1})
 
         # Check result, each segment is 1 second, with duration=1 only one segment should be returned
         expectedResult = b''.join(streams[1:2])
         self.assertEqual(streamlinkResult, expectedResult)
 
-    def test_hls_encryted_aes128(self):
+    def test_hls_encrypted_aes128(self):
         # Encryption parameters
         aesKey = os.urandom(16)
         aesIv = os.urandom(16)
@@ -122,6 +123,38 @@ class TestHLS(unittest.TestCase):
         expectedResult = b''.join(clearStreams[1:] + clearStreams)
         self.assertEqual(streamlinkResult, expectedResult)
 
+    def test_hls_encrypted_aes128_key_uri_override(self):
+        aesKey = os.urandom(16)
+        aesIv = os.urandom(16)
+        aesKeyInvalid = bytes([ord(aesKey[i:i + 1]) ^ 0xFF for i in range(16)])
+        clearStreams = [os.urandom(1024) for i in range(4)]
+        encryptedStreams = [encrypt(clearStream, aesKey, aesIv) for clearStream in clearStreams]
+
+        masterPlaylist = self.getMasterPlaylist()
+        playlist1 = self.getPlaylist(aesIv, "stream{0}.ts.enc")
+        playlist2 = self.getPlaylist(aesIv, "stream2_{0}.ts.enc") + "#EXT-X-ENDLIST\n"
+
+        mocked_key_uri = None
+        streamlinkResult = None
+        with requests_mock.Mocker() as mock:
+            mock.get("http://mocked/path/master.m3u8", text=masterPlaylist)
+            mock.get("http://mocked/path/playlist.m3u8", [{'text': playlist1}, {'text': playlist2}])
+            mock.get("http://mocked/path/encryption_key.key", content=aesKeyInvalid)
+            mocked_key_uri = mock.get("http://real-mocked/path/encryption_key.key", content=aesKey)
+            for i, encryptedStream in enumerate(encryptedStreams):
+                mock.get("http://mocked/path/stream{0}.ts.enc".format(i), content=encryptedStream)
+            for i, encryptedStream in enumerate(encryptedStreams):
+                mock.get("http://mocked/path/stream2_{0}.ts.enc".format(i), content=encryptedStream)
+
+            streamlinkResult = self.start_streamlink(
+                "http://mocked/path/master.m3u8",
+                hls_segment_key_uri="{scheme}://real-{netloc}{path}{query}"
+            )
+
+        self.assertTrue(mocked_key_uri.called)
+        expectedResult = b''.join(clearStreams[1:] + clearStreams)
+        self.assertEqual(streamlinkResult, expectedResult)
+
 
 @patch('streamlink.stream.hls.FFMPEGMuxer.is_usable', Mock(return_value=True))
 class TestHlsExtAudio(unittest.TestCase):
```
tarkah: That's what I was thinking of last night, but was in the same boat about thinking it was probably too much. Though, it does the job. I removed the `aesKeyInvalid` line and mock since that mock is never called in the test anyway. Test works in Py2 & 3. Thanks for all your help!
tarkah: Damn, I didn't think through those scenarios. Sorry, I don't have much experience with writing tests, so I appreciate your helpful explanations of everything!

This feels like a lot to properly isolate and test "using the right key URI" and "does our template string work". Would it be better to create a new function in the module, like the following:

```py
def key_uri(key, key_uri_override=None):
    if key_uri_override:
        p = urlparse(key.uri)
        key_uri = LazyFormatter.format(
            self.key_uri_override,
            url=key.uri,
            scheme=p.scheme,
            netloc=p.netloc,
            path=p.path,
            query=p.query,
        )
    else:
        key_uri = key.uri

    return key_uri
```

And just test this against our hls-segment-key-uri templated string?
tarkah: @back-to I've removed the unnecessary bracket escaping / format

Patch:
diff --git a/src/streamlink/stream/hls.py b/src/streamlink/stream/hls.py
index ef9d0312..ed1f93eb 100644
--- a/src/streamlink/stream/hls.py
+++ b/src/streamlink/stream/hls.py
@@ -6,6 +6,7 @@ from collections import defaultdict, namedtuple
 from Crypto.Cipher import AES
 from requests.exceptions import ChunkedEncodingError
 
+from streamlink.compat import urlparse
 from streamlink.exceptions import StreamError
 from streamlink.stream import hls_playlist
 from streamlink.stream.ffmpegmux import FFMPEGMuxer, MuxedStream
@@ -13,6 +14,7 @@ from streamlink.stream.http import HTTPStream
 from streamlink.stream.segmented import (SegmentedStreamReader,
                                          SegmentedStreamWriter,
                                          SegmentedStreamWorker)
+from streamlink.utils import LazyFormatter
 
 log = logging.getLogger(__name__)
 Sequence = namedtuple("Sequence", "num segment")
@@ -64,7 +66,18 @@ class HLSStreamWriter(SegmentedStreamWriter):
         if not self.key_uri_override and not key.uri:
             raise StreamError("Missing URI to decryption key")
 
-        key_uri = self.key_uri_override if self.key_uri_override else key.uri
+        if self.key_uri_override:
+            p = urlparse(key.uri)
+            key_uri = LazyFormatter.format(
+                self.key_uri_override,
+                url=key.uri,
+                scheme=p.scheme,
+                netloc=p.netloc,
+                path=p.path,
+                query=p.query,
+            )
+        else:
+            key_uri = key.uri
 
         if self.key_uri != key_uri:
             res = self.session.http.get(key_uri, exception=StreamError,
diff --git a/src/streamlink_cli/argparser.py b/src/streamlink_cli/argparser.py
index 40ea303b..8861ceba 100644
--- a/src/streamlink_cli/argparser.py
+++ b/src/streamlink_cli/argparser.py
@@ -817,7 +817,16 @@ def build_parser():
         URI to segment encryption key. If no URI is specified, the URI contained
         in the segments will be used.
 
-        Example: --hls-segment-key-uri "https://example.com/hls/encryption_key"
+        URI can be templated using the following variables, which will be
+        replaced with its respective part from the source segment URI:
+
+          {url} {scheme} {netloc} {path} {query}
+
+        Examples:
+
+          --hls-segment-key-uri "https://example.com/hls/encryption_key"
+          --hls-segment-key-uri "{scheme}://1.2.3.4{path}{query}"
+          --hls-segment-key-uri "{scheme}://{netloc}/custom/path/to/key"
 
         Default is None.
         """


Test patch:
diff --git a/tests/streams/test_hls.py b/tests/streams/test_hls.py
index 61db631e..0fc1cd1a 100644
--- a/tests/streams/test_hls.py
+++ b/tests/streams/test_hls.py
@@ -59,13 +59,14 @@ class TestHLS(unittest.TestCase):
 
         return playlist + playlistEnd
 
-    def start_streamlink(self, masterPlaylist, kwargs=None):
+    def start_streamlink(self, masterPlaylist, hls_segment_key_uri=None, kwargs=None):
         kwargs = kwargs or {}
         log.info("Executing streamlink")
         streamlink = Streamlink()
 
         # Set to default value to avoid a test fail if the default change
         streamlink.set_option("hls-live-edge", 3)
+        streamlink.set_option("hls-segment-key-uri", hls_segment_key_uri)
 
         masterStream = hls.HLSStream.parse_variant_playlist(streamlink, masterPlaylist, **kwargs)
         stream = masterStream["1080p (source)"].open()
@@ -86,13 +87,13 @@ class TestHLS(unittest.TestCase):
 
             # Start streamlink on the generated stream
             streamlinkResult = self.start_streamlink("http://mocked/path/master.m3u8",
-                                                     {'start_offset': 1, 'duration': 1})
+                                                     kwargs={'start_offset': 1, 'duration': 1})
 
         # Check result, each segment is 1 second, with duration=1 only one segment should be returned
         expectedResult = b''.join(streams[1:2])
         self.assertEqual(streamlinkResult, expectedResult)
 
-    def test_hls_encryted_aes128(self):
+    def test_hls_encrypted_aes128(self):
         # Encryption parameters
         aesKey = os.urandom(16)
         aesIv = os.urandom(16)
@@ -122,6 +123,40 @@ class TestHLS(unittest.TestCase):
         expectedResult = b''.join(clearStreams[1:] + clearStreams)
         self.assertEqual(streamlinkResult, expectedResult)
 
+    def test_hls_encrypted_aes128_key_uri_override(self):
+        aesKey = os.urandom(16)
+        aesIv = os.urandom(16)
+        aesKeyInvalid = bytes([ord(aesKey[i:i + 1]) ^ 0xFF for i in range(16)])
+        clearStreams = [os.urandom(1024) for i in range(4)]
+        encryptedStreams = [encrypt(clearStream, aesKey, aesIv) for clearStream in clearStreams]
+
+        masterPlaylist = self.getMasterPlaylist()
+        playlist1 = self.getPlaylist(aesIv, "stream{0}.ts.enc")
+        playlist2 = self.getPlaylist(aesIv, "stream2_{0}.ts.enc") + "#EXT-X-ENDLIST\n"
+
+        mocked_key_uri_default = None
+        mocked_key_uri_override = None
+        streamlinkResult = None
+        with requests_mock.Mocker() as mock:
+            mock.get("http://mocked/path/master.m3u8", text=masterPlaylist)
+            mock.get("http://mocked/path/playlist.m3u8", [{'text': playlist1}, {'text': playlist2}])
+            mocked_key_uri_default = mock.get("http://mocked/path/encryption_key.key", content=aesKeyInvalid)
+            mocked_key_uri_override = mock.get("http://real-mocked/path/encryption_key.key", content=aesKey)
+            for i, encryptedStream in enumerate(encryptedStreams):
+                mock.get("http://mocked/path/stream{0}.ts.enc".format(i), content=encryptedStream)
+            for i, encryptedStream in enumerate(encryptedStreams):
+                mock.get("http://mocked/path/stream2_{0}.ts.enc".format(i), content=encryptedStream)
+
+            streamlinkResult = self.start_streamlink(
+                "http://mocked/path/master.m3u8",
+                hls_segment_key_uri="{scheme}://real-{netloc}{path}{query}"
+            )
+
+        self.assertFalse(mocked_key_uri_default.called)
+        self.assertTrue(mocked_key_uri_override.called)
+        expectedResult = b''.join(clearStreams[1:] + clearStreams)
+        self.assertEqual(streamlinkResult, expectedResult)
+
 
 @patch('streamlink.stream.hls.FFMPEGMuxer.is_usable', Mock(return_value=True))
 class TestHlsExtAudio(unittest.TestCase):