backend improvements

backend/main.py

@@ -277,7 +277,7 @@ def stream_video():
 
         return forwarded_headers
 
-    def proxy_response(target_url, content_type_override=None, referer_override=None):
+    def proxy_response(target_url, content_type_override=None, referer_override=None, upstream_headers=None):
         # Extract the base domain to spoof the referer
         request_referer = request.args.get('referer')
         if referer_override:
@@ -290,6 +290,10 @@ def stream_video():
         dbg(f"proxy_response target={target_url} referer={referer}")
 
         safe_request_headers = build_upstream_headers(referer)
+        if isinstance(upstream_headers, dict):
+            for key, value in upstream_headers.items():
+                if value:
+                    safe_request_headers[key] = value
 
         # Pass through Range headers so the browser can 'sniff' the video
         if 'Range' in request.headers:
@@ -315,9 +319,12 @@ def stream_video():
                 encoding = resp.encoding
                 resp.close()
                 dbg("detected m3u8 by content sniff")
+                upstream_for_playlist = dict(safe_request_headers)
+                upstream_for_playlist.pop('Range', None)
                 return proxy_hls_playlist(
                     target_url,
                     referer_hint=referer,
+                    upstream_headers=upstream_for_playlist,
                     prefetched_body=body_bytes,
                     prefetched_base_url=base_url,
                     prefetched_encoding=encoding,
@@ -378,20 +385,39 @@ def stream_video():
         body = "\n".join(rewritten)
         return Response(body, status=200, content_type='application/vnd.apple.mpegurl')
 
-    def proxy_hls_playlist(playlist_url, referer_hint=None, prefetched_body=None, prefetched_base_url=None, prefetched_encoding=None):
+    def proxy_hls_playlist(playlist_url, referer_hint=None, prefetched_body=None, prefetched_base_url=None, prefetched_encoding=None, upstream_headers=None):
         dbg(f"proxy_hls_playlist url={playlist_url} referer_hint={referer_hint}")
         base_url = prefetched_base_url or playlist_url
         body_text = None
         if prefetched_body is None:
             headers = build_upstream_headers(referer_hint or "")
+            if isinstance(upstream_headers, dict):
+                for key, value in upstream_headers.items():
+                    if value:
+                        headers[key] = value
             if 'User-Agent' not in headers:
                 headers['User-Agent'] = 'Mozilla/5.0'
             if 'Accept' not in headers:
                 headers['Accept'] = '*/*'
             resp = session.get(playlist_url, headers=headers, stream=True, timeout=30)
-            resp.raise_for_status()
             base_url = resp.url
 
+            if resp.status_code >= 400:
+                forwarded_headers = build_forwarded_headers(resp, target_url=base_url)
+                if request.method == 'HEAD':
+                    resp.close()
+                    return Response("", status=resp.status_code, headers=forwarded_headers)
+
+                def generate():
+                    try:
+                        for chunk in resp.iter_content(chunk_size=1024 * 16):
+                            if chunk:
+                                yield chunk
+                    finally:
+                        resp.close()
+
+                return Response(generate(), status=resp.status_code, headers=forwarded_headers)
+
             if request.method == 'HEAD':
                 forwarded_headers = build_forwarded_headers(resp, target_url=base_url)
                 resp.close()
@@ -459,6 +485,11 @@ def stream_video():
         except Exception as e:
             return jsonify({"error": str(e)}), 500
 
+    def extract_referer(headers):
+        if not isinstance(headers, dict):
+            return None
+        return headers.get('Referer') or headers.get('referer')
+
     try:
         # Configure yt-dlp options
         ydl_opts = {
@@ -488,21 +519,41 @@ def stream_video():
             # Try to get the URL from the info dict (works for progressive downloads)
             stream_url = info.get('url')
             protocol = info.get('protocol')
-            
+            selected_format = None
+
             # If no direct URL, try to get it from formats
-            if not stream_url and 'formats' in info:
+            if 'formats' in info:
-                # Find the best format that has a URL
+                if info.get('format_id'):
-                for fmt in info['formats']:
+                    for fmt in info['formats']:
-                    if fmt.get('url'):
+                        if fmt.get('format_id') == info.get('format_id'):
-                        stream_url = fmt.get('url')
+                            selected_format = fmt
-                        break
+                            break
-            
+                if not selected_format and stream_url:
+                    for fmt in info['formats']:
+                        if fmt.get('url') == stream_url:
+                            selected_format = fmt
+                            break
+                if not selected_format:
+                    for fmt in info['formats']:
+                        if fmt.get('url'):
+                            selected_format = fmt
+                            break
+
+                if not stream_url and selected_format:
+                    stream_url = selected_format.get('url')
+
             if not stream_url:
                 return jsonify({"error": "Could not extract stream URL"}), 500
 
+            upstream_headers = None
+            if selected_format and isinstance(selected_format.get('http_headers'), dict):
+                upstream_headers = selected_format['http_headers']
+            elif isinstance(info.get('http_headers'), dict):
+                upstream_headers = info['http_headers']
+
             referer_hint = None
-            if isinstance(info.get('http_headers'), dict):
+            if upstream_headers:
-                referer_hint = info['http_headers'].get('Referer') or info['http_headers'].get('referer')
+                referer_hint = extract_referer(upstream_headers)
             if not referer_hint:
                 parsed = urllib.parse.urlparse(video_url)
                 referer_hint = f"{parsed.scheme}://{parsed.netloc}/"
@@ -510,18 +561,18 @@ def stream_video():
 
             if protocol and 'm3u8' in protocol:
                 dbg("protocol indicates hls")
-                return proxy_hls_playlist(stream_url, referer_hint)
+                return proxy_hls_playlist(stream_url, referer_hint, upstream_headers=upstream_headers)
 
             if is_hls(stream_url):
                 dbg("stream_url is hls")
-                return proxy_hls_playlist(stream_url, referer_hint)
+                return proxy_hls_playlist(stream_url, referer_hint, upstream_headers=upstream_headers)
 
             if is_dash(stream_url):
                 dbg("stream_url is dash")
-                return proxy_response(stream_url, content_type_override='application/dash+xml', referer_override=referer_hint)
+                return proxy_response(stream_url, content_type_override='application/dash+xml', referer_override=referer_hint, upstream_headers=upstream_headers)
 
             dbg("stream_url is direct media")
-            return proxy_response(stream_url, referer_override=referer_hint)
+            return proxy_response(stream_url, referer_override=referer_hint, upstream_headers=upstream_headers)
     except Exception as e:
         return jsonify({"error": str(e)}), 500