From be0a868861dfab53963029ba52576518138680ad Mon Sep 17 00:00:00 2001
From: Roy Olav Purser <admin@royolav.net>
Date: Tue, 11 May 2021 14:00:43 +0200
Subject: [PATCH] add proxy object

---
 stream.py | 82 ++++++++++++++++++++++++++++---------------------------
 1 file changed, 42 insertions(+), 40 deletions(-)

diff --git a/stream.py b/stream.py
index 5d5ec33..743c315 100755
--- a/stream.py
+++ b/stream.py
@@ -11,23 +11,25 @@ import re
 providers = {}
 providers["nrk"] = "https://tv.nrk.no"
 providers["svt"] = "https://svtplay.se"
-providers["youtube"] = "https://youtu.be"
+providers["youtube"] = "https://www.youtube.com/watch?v="
 providers["twitch"] = "https://twitch.tv"
 
-proxies = {}
-proxies_raw = {}
-proxies_req = {}
+class ProxyElem():
+    def __init__(self, proxy):
+        self.stream = streamlink.Streamlink()
+        self.proxy = proxy
+        self.req = {}
+        self.stream.set_option("http-timeout", 2.0)
+        if proxy is not None:
+            self.stream.set_option("https-proxy", "socks5://" + proxy)
+            self.stream.set_option("http-proxy", "socks5://" + proxy)
+            self.req["http"] = "socks5://" + proxy
+            self.req["https"] = "socks5://" + proxy
 
+proxies = {}
 for key in providers:
-    proxies[key] = streamlink.Streamlink()
     proxy = os.environ.get(key.upper() + "_PROXY")
-    proxies[key].set_option("http-timeout", 2.0)
-    if proxy is not None:
-        proxies[key].set_option("https-proxy", "socks5://" + proxy)
-        proxies_raw[key] = proxy
-        proxies_req[key] = {}
-        proxies_req[key]["http"] = "socks5://" + proxy
-        proxies_req[key]["https"] = "socks5://" + proxy
+    proxies[key] = ProxyElem(proxy)
 
 playlist = None
 icecast_server = os.environ.get("ICECAST_SERVER")
@@ -74,19 +76,17 @@ def get_proxy_url(proxy, current, path):
     return presp.text
 
 def rewrite(current, provider):
-    proxy_req = proxies_req.get(provider)
-    proxy = proxies_raw.get(provider)
-    resp = requests.head(current, proxies=proxy_req)
+    proxy = proxies.get(provider)
+    resp = requests.head(current, proxies=proxy.req)
     ctype = resp.headers.get("Content-Type")
     if ctype is None:
         return None
     else:
         if "mpegurl" not in ctype.lower():
             return None
-    resp = requests.get(current, proxies=proxy_req)
+    resp = requests.get(current, proxies=proxy.req)
     ndata = None
     if resp.text is not None:
-        ndata = ""
         links = []
         for line in resp.text.splitlines():
             if line.startswith("#EXT-X-KEY:METHOD="):
@@ -94,32 +94,34 @@ def rewrite(current, provider):
                 if len(matches) == 1:
                     ldata = {}
                     ldata["upstream"] = urllib.parse.urljoin(current, matches[0])
-                    ldata["proxy"] = proxy
-                    ldata["proxied"] = isinstance(proxy, str)
+                    ldata["proxy"] = proxy.proxy
+                    ldata["proxied"] = isinstance(proxy.proxy, str)
                     links.append(ldata)
             elif line.startswith("#"):
                 pass
             else:
                 ldata = {}
                 ldata["upstream"] = urllib.parse.urljoin(current, line)
-                ldata["proxy"] = proxy
-                ldata["proxied"] = isinstance(proxy, str)
+                ldata["proxy"] = proxy.proxy
+                ldata["proxied"] = isinstance(proxy.proxy, str)
                 links.append(ldata)
-        presp = requests.post(proxy_server, json=links)
-        if isinstance(presp.text, str):
-            print(presp.text)
-            links = json.loads(presp.text)
-        for line in resp.text.splitlines():
-            if line.startswith("#EXT-X-KEY:METHOD="):
-                matches = re.findall(r'(?<=URI=").+(?=")', line)
-                if len(matches) == 1:
-                    new_url = links.pop(0)
-                    ndata += re.sub(r'URI=".+"', f'URI="{new_url}"', line)
-            elif line.startswith("#"):
-                ndata += line
-            else:
-                ndata += links.pop(0)
-            ndata += "\n"
+        if isinstance(proxy_server, str):
+            ndata = ""
+            presp = requests.post(proxy_server, json=links)
+            if isinstance(presp.text, str):
+                print(presp.text)
+                links = json.loads(presp.text)
+            for line in resp.text.splitlines():
+                if line.startswith("#EXT-X-KEY:METHOD="):
+                    matches = re.findall(r'(?<=URI=").+(?=")', line)
+                    if len(matches) == 1:
+                        new_url = links.pop(0)
+                        ndata += re.sub(r'URI=".+"', f'URI="{new_url}"', line)
+                elif line.startswith("#"):
+                    ndata += line
+                else:
+                    ndata += links.pop(0)
+                ndata += "\n"
         return ndata
 
 class MainHandler(tornado.web.RequestHandler):
@@ -153,19 +155,19 @@ class MainHandler(tornado.web.RequestHandler):
         proxy = None
         if provider in providers.keys():
             proxy = proxies.get(provider)
+            path = self.request.path
+            if provider == "youtube":
+                path = path.strip("/")
             src = providers[provider] + self.request.path
             try:
                 resp = requests.head(src, allow_redirects=True)
                 if resp is not None:
                     src = resp.url
-                    if isinstance(src, str) and "consent.youtube.com" in src:
-                        video_id = self.request.path.strip("/")
-                        src = f'https://www.youtube.com/watch?v={video_id}'
             except Exception as e:
                 print(e)
             else:
                 try:
-                    streams = proxy.streams(src)
+                    streams = proxy.stream.streams(src)
                     for key in reversed(streams):
                         stream = streams.get(key)
                         if hasattr(stream, "url"):