add fallback to youtube-dl

backend/install.sh

@@ -2,4 +2,4 @@
 virtualenv --python=$(which python3) /app/venv
 source /app/venv/bin/activate
 pip install --upgrade pip
-pip install --upgrade streamlink tornado aiohttp aiohttp-socks
+pip install --upgrade streamlink youtube-dl tornado aiohttp aiohttp-socks

backend/stream.py

@@ -8,11 +8,11 @@ import time
 import base64
 import logging
 import asyncio
-import streamlink
 import tornado.web
 import tornado.routing
 import aiohttp
 import aiohttp_socks
+import stream_providers
 
 logging.basicConfig(format='[%(filename)s:%(lineno)d] %(message)s', stream=sys.stdout, level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -30,13 +30,6 @@ proxy_server = os.environ.get("PROXY_SERVER")
 class ProxyElem():
     def __init__(self, proxy):
         self.proxy = proxy
-    def stream(self):
-        session = streamlink.Streamlink()
-        session.set_option("http-timeout", 2.0)
-        if self.proxy is not None:
-            session.set_option("https-proxy", "socks5://" + self.proxy)
-            session.set_option("http-proxy", "socks5://" + self.proxy)
-        return session
     def local(self):
         timeout = aiohttp.ClientTimeout(total=1)
         return aiohttp.ClientSession(timeout=timeout)
@@ -99,15 +92,6 @@ class AsyncSession():
         resp = await self.sdata.resp
         return AsyncSessionData(resp, self.sdata.current)
 
-class StreamRunner():
-    def __init__(self, upstream, proxy):
-        self.upstream = upstream
-        self.proxy = proxy
-    def stream(self):
-        return self.proxy.stream().streams(self.upstream)
-    async def run(self):
-        return await asyncio.to_thread(self.stream)
-
 proxies = {}
 for key in providers:
     proxies[key] = []
@@ -365,23 +349,14 @@ class MainHandler(tornado.web.RequestHandler):
     async def handle_stream(self, handler, redir):
         upstream = None
         streams = None
-        runner = StreamRunner(handler.upstream, handler.proxy)
-        start_time = time.time_ns()
-        for i in range(5):
-            try:
-                streams = await runner.run()
-            except Exception as e:
-                logger.info(e)
-            else:
-                break
-        stop_time = time.time_ns()
-        logger.info((stop_time - start_time) // 1_000_000)
-        if streams is not None:
-            for key in reversed(streams):
-                stream = streams.get(key)
-                if hasattr(stream, "url"):
-                    upstream = stream.url
-                    break
+        youtube_stream_future = stream_providers.ytdl_get(handler.upstream, handler.proxy, logger)
+        streamlink_stream_future = stream_providers.streamlink_get(handler.upstream, handler.proxy, logger)
+        youtube_stream = await youtube_stream_future
+        streamlink_stream = await streamlink_stream_future
+        if isinstance(streamlink_stream, str):
+            upstream = streamlink_stream
+        elif isinstance(youtube_stream, str):
+            upstream = youtube_stream
         if upstream is None:
             logger.info(f'invalid upstream ({handler.provider})')
             self.set_status(404)

backend/stream_providers.py

@@ -0,0 +1,88 @@
+#!/usr/bin/env python3
+import youtube_dl
+import streamlink
+import asyncio
+
+class DummyLogger():
+    def debug(self, msg):
+        pass
+    def warning(self, msg):
+        pass
+    def error(self, msg):
+        pass
+
+class StreamProvider():
+    def __init__(self, upstream, proxy):
+        self.upstream = upstream
+        self.proxy = None
+        proxy = str(proxy)
+        if len(proxy) > 5:
+            self.proxy = "socks5://" + proxy
+
+class StreamlinkRunner(StreamProvider):
+    def stream(self):
+        session = streamlink.Streamlink()
+        session.set_option("http-timeout", 2.0)
+        if self.proxy is not None:
+            session.set_option("https-proxy", self.proxy)
+            session.set_option("http-proxy", self.proxy)
+        streams = session.streams(self.upstream)
+        if streams is not None:
+            for key in reversed(streams):
+                stream = streams.get(key)
+                if hasattr(stream, "url"):
+                    return stream.url
+        return None
+    async def run(self):
+        return await asyncio.to_thread(self.stream)
+
+class YoutubeRunner(StreamProvider):
+    def stream(self):
+        opts = {}
+        opts["logger"] = DummyLogger()
+        if isinstance(self.proxy, str):
+            opts["proxy"] = self.proxy
+        best_url = None
+        with youtube_dl.YoutubeDL(opts) as ydl:
+            info = ydl.extract_info(self.upstream, download=False)
+            vformats = info.get("formats")
+            best_format = {}
+            best_format["quality"] = 0
+            if isinstance(vformats, list):
+                for vformat in vformats:
+                    acodec = vformat.get("acodec")
+                    vcodec = vformat.get("vcodec")
+                    new_quality = vformat.get("quality")
+                    old_quality = best_format.get("quality")
+                    new_url = vformat.get("url")
+                    if (isinstance(acodec, str) and
+                            isinstance(vcodec, str) and
+                            isinstance(new_quality, int) and
+                            isinstance(old_quality, int) and
+                            isinstance(new_url, str) and
+                            acodec != "none" and
+                            vcodec != "none" and
+                            new_quality > old_quality):
+                        best_format = vformat
+                        best_url = new_url
+        return best_url
+    async def run(self):
+        return await asyncio.to_thread(self.stream)
+
+async def ytdl_get(upstream, proxy, logger):
+    url = None
+    try:
+        runner = YoutubeRunner(upstream, proxy)
+        url = await runner.run()
+    except Exception as e:
+        logger.info(e)
+    return url
+ 
+async def streamlink_get(upstream, proxy, logger):
+    url = None
+    try:
+        runner = StreamlinkRunner(upstream, proxy)
+        url = await runner.run()
+    except Exception as e:
+        logger.info(e)
+    return url