From 4267f6a7ca9a4f341438efadb22f8101ad200338 Mon Sep 17 00:00:00 2001
From: Roy Olav Purser <admin@royolav.net>
Date: Sun, 13 Feb 2022 14:15:33 +0100
Subject: [PATCH] add download support

---
 backend/config.py  |  3 +++
 backend/stream.py  | 40 ++++++++++++++++++++++++++++++++--------
 frontend/script.js |  9 ++++-----
 3 files changed, 39 insertions(+), 13 deletions(-)

diff --git a/backend/config.py b/backend/config.py
index 5b8b023..5011115 100644
--- a/backend/config.py
+++ b/backend/config.py
@@ -45,11 +45,13 @@ proxy_server = os.environ.get("PROXY_SERVER")
 @dataclasses.dataclass
 class LinkWithType:
     upstream: str
+    download: bool = False
     ctype: None | str = None
 
 
 class ProxyCreateLink(pydantic.BaseModel):
     upstream: None | pydantic.HttpUrl
+    download: None | pydantic.StrictBool
     ctype: None | pydantic.StrictStr
     region: None | pydantic.StrictStr
 
@@ -91,6 +93,7 @@ class ProxyElem:
                         {
                             "upstream": url.upstream,
                             "ctype": url.ctype,
+                            "download": url.download,
                             "region": self.region,
                         }
                     )
diff --git a/backend/stream.py b/backend/stream.py
index a73039b..e9bb196 100755
--- a/backend/stream.py
+++ b/backend/stream.py
@@ -2,13 +2,14 @@
 import json
 import sys
 import base64
+import dataclasses
+from typing import Optional, cast, Any
 import logging
 import asyncio
 import tornado.web
 import tornado.routing
 import config
 import stream_providers
-from typing import Optional, cast, Any
 
 logging.basicConfig(
     format="[%(filename)s:%(lineno)d] %(message)s",
@@ -85,6 +86,14 @@ class UpstreamHandler:
                             future.cancel()
 
 
+@dataclasses.dataclass
+class VideoInfo:
+    upstream: None | str = None
+    download: None | str = None
+    poster: None | str = None
+    ctype: None | str = None
+
+
 class NoDataError(Exception):
     pass
 
@@ -127,10 +136,11 @@ class MainHandler(tornado.web.RequestHandler):
         if provider_data is None:
             raise NoDataError()
 
-        video_info = {}
+        video_info = VideoInfo()
+        video_info.ctype = provider_data.ctype()
         if handler.direct:
-            video_info["upstream"] = provider_data.upstream()
-            video_info["poster"] = provider_data.thumbnail()
+            video_info.upstream = provider_data.upstream()
+            video_info.poster = provider_data.thumbnail()
         else:
             proxied = await handler.proxy.proxy_url(
                 [
@@ -138,16 +148,30 @@ class MainHandler(tornado.web.RequestHandler):
                         upstream=provider_data.upstream(),
                         ctype=provider_data.proxy_ctype(),
                     ),
+                    config.LinkWithType(
+                        upstream=provider_data.upstream(),
+                        ctype=provider_data.proxy_ctype(),
+                        download=True,
+                    ),
                     config.LinkWithType(
                         upstream=provider_data.thumbnail(),
                     ),
                 ]
             )
-            video_info["upstream"] = proxied[0]
-            video_info["poster"] = proxied[1]
+            video_info.upstream = proxied[0]
+            video_info.download = proxied[1]
+            if (
+                isinstance(video_info.ctype, str)
+                and "mpegurl" in video_info.ctype.lower()
+            ):
+                video_info.download = None
+            video_info.poster = proxied[2]
 
-        video_info["ctype"] = provider_data.ctype()
-        return (video_info, provider_data.meta(), provider_data.title())
+        return (
+            dataclasses.asdict(video_info),
+            provider_data.meta(),
+            provider_data.title(),
+        )
 
     async def handle_raw(self, handler):
         try:
diff --git a/frontend/script.js b/frontend/script.js
index 444f071..6716c90 100644
--- a/frontend/script.js
+++ b/frontend/script.js
@@ -15,7 +15,7 @@
         options.plugins.chromecast = {};
         options.plugins.chromecast.addButtonToControlBar = false;
         const player = videojs(video, options);
-        if((info.poster instanceof String) || ((typeof info.poster) == "string")) {
+        if(info.poster) {
             player.poster(info.poster);
         }
         const source = {};
@@ -24,8 +24,7 @@
         player.src(source);
 
         const dl = document.createElement("a");
-        dl.href = info.upstream;
-        dl.setAttribute("download", true);
+        dl.href = info.download;
 
         const canPlayTypeRaw = player.canPlayType(info.ctype);
         const canPlayType = (canPlayTypeRaw === "maybe") || (canPlayTypeRaw === "probably");
@@ -66,7 +65,7 @@
             bigCastButton.classList.add("vjs-big-chromecast");
             player.addChild(bigCastButtonWrapper);
  
-            if(!info.ctype.toLowerCase().includes("mpegurl")) {
+            if(info.download) {
                 bigDownloadButton.classList.add("vjs-control");
                 bigDownloadButton.classList.add("vjs-button");
                 bigDownloadButton.classList.add("fas");
@@ -77,7 +76,7 @@
 
             player.one("error", () => {
                 player.removeChild(bigCastButtonWrapper);
-                if(!info.ctype.toLowerCase().includes("mpegurl")) {
+                if(info.download) {
                     player.removeChild(bigDownloadButtonWrapper);
                 }
             });