tubearchivist/tubearchivist/home/src/download/yt_dlp_base.py

"""
functionality:
- base class to make all calls to yt-dlp
- handle yt-dlp errors
"""

import os
from http import cookiejar
from io import StringIO

import yt_dlp
from home.src.ta.ta_redis import RedisArchivist


class YtWrap:
    """wrap calls to yt"""

    OBS_BASE = {
        "default_search": "ytsearch",
        "quiet": True,
        "check_formats": "selected",
        "socket_timeout": 2,
    }

    def __init__(self, obs_request, config=False):
        self.obs_request = obs_request
        self.config = config
        self.build_obs()

    def build_obs(self):
        """build yt-dlp obs"""
        self.obs = self.OBS_BASE.copy()
        self.obs.update(self.obs_request)
        if self.config:
            self.add_cookie()

    def add_cookie(self):
        """add cookie if enabled"""
        if self.config["downloads"]["cookie_import"]:
            cookie_io = CookieHandler(self.config).get()
            self.obs["cookiefile"] = cookie_io

    def download(self, url):
        """make download request"""
        with yt_dlp.YoutubeDL(self.obs) as ydl:
            try:
                ydl.download([url])
            except yt_dlp.utils.DownloadError:
                print(f"{url}: failed to download.")
                return False

        return True

    def extract(self, url):
        """make extract request"""
        try:
            response = yt_dlp.YoutubeDL(self.obs).extract_info(url)
        except cookiejar.LoadError:
            print("cookie file is invalid")
            return False
        except (yt_dlp.utils.ExtractorError, yt_dlp.utils.DownloadError):
            print(f"{url}: failed to get info from youtube")
            return False

        return response


class CookieHandler:
    """handle youtube cookie for yt-dlp"""

    def __init__(self, config):
        self.cookie_io = False
        self.config = config

    def get(self):
        """get cookie io stream"""
        cookie = RedisArchivist().get_message("cookie")
        self.cookie_io = StringIO(cookie)
        return self.cookie_io

    def import_cookie(self):
        """import cookie from file"""
        cache_path = self.config["application"]["cache_dir"]
        import_path = os.path.join(cache_path, "import", "cookies.google.txt")
        with open(import_path, encoding="utf-8") as cookie_file:
            cookie = cookie_file.read()

        self.set_cookie(cookie)

        os.remove(import_path)
        print("cookie: import successful")

    def set_cookie(self, cookie):
        """set cookie str and activate in cofig"""
        RedisArchivist().set_message("cookie", cookie)
        path = ".downloads.cookie_import"
        RedisArchivist().set_message("config", True, path=path)
        self.config["downloads"]["cookie_import"] = True
        print("cookie: activated and stored in Redis")

    @staticmethod
    def revoke():
        """revoke cookie"""
        RedisArchivist().del_message("cookie")
        RedisArchivist().set_message(
            "config", False, path=".downloads.cookie_import"
        )
        print("cookie: revoked")

    def validate(self):
        """validate cookie using the liked videos playlist"""
        print("validating cookie")
        obs_request = {
            "skip_download": True,
            "extract_flat": True,
        }
        validator = YtWrap(obs_request, self.config)
        response = validator.extract("LL")

        # update in redis to avoid expiring
        modified = validator.obs["cookiefile"].getvalue()
        if modified:
            RedisArchivist().set_message("cookie", modified)

        if not response:
            mess_dict = {
                "status": "message:download",
                "level": "error",
                "title": "Cookie validation failed, exiting...",
                "message": "",
            }
            RedisArchivist().set_message(
                "message:download", mess_dict, expire=4
            )
            print("cookie validation failed, exiting...")

        return bool(response)