From 2c0b0afb391a5d8c8f678533cad18019fce94c36 Mon Sep 17 00:00:00 2001
From: neoarz <tyrantneo740@gmail.com>
Date: Fri, 3 Oct 2025 12:49:50 -0400
Subject: [PATCH] fix(download): support for yt-dlp cookies in media download

Introduces environment variable support for specifying yt-dlp cookies via file path or direct content for the media download command. Updates error handling to provide clearer messages when cookies are required and ensures temporary cookie files are cleaned up. Also updates .env.example and .gitignore to reflect these changes.
---
 .env.example           |  5 +++-
 .gitignore             |  2 ++
 cogs/media/download.py | 53 +++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/.env.example b/.env.example
index 49b3f2b..cafd098 100644
--- a/.env.example
+++ b/.env.example
@@ -3,4 +3,7 @@ PREFIX=YOUR_BOT_PREFIX_HERE
 INVITE_LINK=YOUR_BOT_INVITE_LINK_HERE
 
 # Commands you want to disable (comma separated)
-DISABLED_COGS=general.context_menus 
+DISABLED_COGS=general.context_menus
+
+# yt-dlp cookies (for media/download command)
+YTDLP_COOKIE_FILE=/absolute/path/to/cookies.txt
diff --git a/.gitignore b/.gitignore
index 00170c9..9f0c101 100644
--- a/.gitignore
+++ b/.gitignore
@@ -149,3 +149,5 @@ cython_debug/
 # Log file
 discord.log
 
+cogs/media/files/
+
diff --git a/cogs/media/download.py b/cogs/media/download.py
index cad1e66..79fbaaf 100644
--- a/cogs/media/download.py
+++ b/cogs/media/download.py
@@ -89,6 +89,7 @@ def download_command():
             processing_msg = await context.send(embed=processing_embed)
 
         temp_dir = tempfile.mkdtemp()
+        temp_cookie_file = None
         
         ydl_opts = {
             'format': 'bestvideo[filesize<200M]+bestaudio[filesize<200M]/best[filesize<200M]/bestvideo+bestaudio/best',
@@ -102,6 +103,45 @@ def download_command():
             'merge_output_format': 'mp4',
         }
 
+        cookie_file_env = os.getenv('YTDLP_COOKIE_FILE') or os.getenv('YT_DLP_COOKIE_FILE')
+        cookies_text_env = os.getenv('YTDLP_COOKIES') or os.getenv('YT_DLP_COOKIES')
+
+        resolved_cookie_path = None
+        if cookie_file_env:
+            resolved_cookie_path = cookie_file_env
+        elif cookies_text_env:
+            try:
+                fd, temp_cookie_file = tempfile.mkstemp(prefix='yt_cookies_', text=True)
+                with os.fdopen(fd, 'w') as tmpf:
+                    tmpf.write(cookies_text_env)
+                resolved_cookie_path = temp_cookie_file
+            except Exception:
+                temp_cookie_file = None
+        else:
+            default_local_cookie = os.path.join(os.path.dirname(__file__), 'files', 'cookies.txt')
+            resolved_cookie_path = default_local_cookie
+
+        if not (resolved_cookie_path and os.path.exists(resolved_cookie_path)):
+            embed = discord.Embed(
+                title="Error",
+                description=(
+                    "Cookies file not found. Provide one via `YTDLP_COOKIE_FILE` or place a file at "
+                    "`cogs/media/files/cookies.txt`."
+                ),
+                color=0xE02B2B,
+            )
+            embed.set_author(name="Media", icon_url="https://yes.nighty.works/raw/y5SEZ9.webp")
+            if interaction is not None:
+                if not interaction.response.is_done():
+                    await interaction.response.send_message(embed=embed, ephemeral=True)
+                else:
+                    await interaction.followup.send(embed=embed, ephemeral=True)
+            else:
+                await context.send(embed=embed, ephemeral=True)
+            return
+
+        ydl_opts['cookiefile'] = resolved_cookie_path
+
         try:
             with yt_dlp.YoutubeDL(ydl_opts) as ydl:
                 info = await asyncio.get_event_loop().run_in_executor(
@@ -232,9 +272,15 @@ def download_command():
                         await context.channel.send(file=file)
                         
         except Exception as e:
+            err_text = str(e)
+            needs_cookies = ('Sign in to confirm' in err_text) or ('cookies' in err_text.lower()) or ('consent' in err_text.lower())
+            if needs_cookies:
+                extra = "\n\nThis source may require authentication. Provide cookies via env: `YTDLP_COOKIE_FILE` (path) or `YTDLP_COOKIES` (contents). See: https://github.com/yt-dlp/yt-dlp/wiki/FAQ#how-do-i-pass-cookies-to-yt-dlp"
+            else:
+                extra = ""
             embed = discord.Embed(
                 title="Error",
-                description=f"Failed to download video: {str(e)}",
+                description=f"Failed to download video: {err_text}{extra}",
                 color=0xE02B2B,
             )
             embed.set_author(name="Media", icon_url="https://yes.nighty.works/raw/y5SEZ9.webp")
@@ -262,5 +308,10 @@ def download_command():
                 os.rmdir(temp_dir)
             except:
                 pass
+            if temp_cookie_file:
+                try:
+                    os.remove(temp_cookie_file)
+                except:
+                    pass
     
     return download
\ No newline at end of file