From afbbc032ee6f1b96256b53c901bfc3259fb571ff Mon Sep 17 00:00:00 2001
From: Roy Olav Purser <admin@royolav.net>
Date: Sun, 9 May 2021 12:37:36 +0200
Subject: [PATCH] add key proxy

---
 stream.py | 27 ++++++++++++++++++---------
 1 file changed, 18 insertions(+), 9 deletions(-)

diff --git a/stream.py b/stream.py
index aaf3a70..60962e3 100755
--- a/stream.py
+++ b/stream.py
@@ -6,6 +6,7 @@ import requests
 import json
 import os
 import urllib.parse
+import re
 
 providers = {}
 providers["nrk"] = "https://tv.nrk.no"
@@ -48,6 +49,16 @@ if icecast_server is not None and stream_server is not None:
                 playlist += f'#EXTINF:0 radio="false", {name}\n'
                 playlist += stream_server + key + "\n"
 
+def get_proxy_url(proxy, current, path):
+    data = {}
+    data["upstream"] = urllib.parse.urljoin(current, path)
+    data["proxy"] = proxy
+    ret = None
+    if proxy is None or proxy_server is None:
+        return data["upstream"]
+    presp = requests.post(proxy_server, json=data)
+    return presp.text
+
 def rewrite(current, provider):
     proxy_req = proxies_req.get(provider)
     proxy = proxies_raw.get(provider)
@@ -64,17 +75,15 @@ def rewrite(current, provider):
     if resp.text is not None:
         ndata = ""
         for line in resp.text.splitlines():
-            if line.startswith("#"):
+            if line.startswith("#EXT-X-KEY:METHOD="):
+                matches = re.findall(r'(?<=URI=").+(?=")', line)
+                if len(matches) == 1:
+                    new_url = get_proxy_url(proxy, current, matches[0])
+                    ndata += re.sub(r'URI=".+"', f'URI="{new_url}"', line)
+            elif line.startswith("#"):
                 ndata += line
             else:
-                data = {}
-                data["upstream"] = urllib.parse.urljoin(current, line)
-                data["proxy"] = proxy
-                if proxy is None or proxy_server is None:
-                    ndata += data["upstream"]
-                else:
-                    presp = requests.post(proxy_server, json=data)
-                    ndata += presp.text
+                ndata += get_proxy_url(proxy, current, line)
             ndata += "\n"
     return ndata