test/test_cookies.py

   1 import unittest
   2 from datetime import datetime, timezone
   3
   4 from yt_dlp import cookies
   5 from yt_dlp.cookies import (
   6     LinuxChromeCookieDecryptor,
   7     MacChromeCookieDecryptor,
   8     WindowsChromeCookieDecryptor,
   9     _get_linux_desktop_environment,
  10     _LinuxDesktopEnvironment,
  11     parse_safari_cookies,
  12     pbkdf2_sha1,
  13 )
  14
  15
  16 class Logger:
  17     def debug(self, message, *args, **kwargs):
  18         print(f'[verbose] {message}')
  19
  20     def info(self, message, *args, **kwargs):
  21         print(message)
  22
  23     def warning(self, message, *args, **kwargs):
  24         self.error(message)
  25
  26     def error(self, message, *args, **kwargs):
  27         raise Exception(message)
  28
  29
  30 class MonkeyPatch:
  31     def __init__(self, module, temporary_values):
  32         self._module = module
  33         self._temporary_values = temporary_values
  34         self._backup_values = {}
  35
  36     def __enter__(self):
  37         for name, temp_value in self._temporary_values.items():
  38             self._backup_values[name] = getattr(self._module, name)
  39             setattr(self._module, name, temp_value)
  40
  41     def __exit__(self, exc_type, exc_val, exc_tb):
  42         for name, backup_value in self._backup_values.items():
  43             setattr(self._module, name, backup_value)
  44
  45
  46 class TestCookies(unittest.TestCase):
  47     def test_get_desktop_environment(self):
  48         """ based on https://chromium.googlesource.com/chromium/src/+/refs/heads/main/base/nix/xdg_util_unittest.cc """
  49         test_cases = [
  50             ({}, _LinuxDesktopEnvironment.OTHER),
  51
  52             ({'DESKTOP_SESSION': 'gnome'}, _LinuxDesktopEnvironment.GNOME),
  53             ({'DESKTOP_SESSION': 'mate'}, _LinuxDesktopEnvironment.GNOME),
  54             ({'DESKTOP_SESSION': 'kde4'}, _LinuxDesktopEnvironment.KDE),
  55             ({'DESKTOP_SESSION': 'kde'}, _LinuxDesktopEnvironment.KDE),
  56             ({'DESKTOP_SESSION': 'xfce'}, _LinuxDesktopEnvironment.XFCE),
  57
  58             ({'GNOME_DESKTOP_SESSION_ID': 1}, _LinuxDesktopEnvironment.GNOME),
  59             ({'KDE_FULL_SESSION': 1}, _LinuxDesktopEnvironment.KDE),
  60
  61             ({'XDG_CURRENT_DESKTOP': 'X-Cinnamon'}, _LinuxDesktopEnvironment.CINNAMON),
  62             ({'XDG_CURRENT_DESKTOP': 'GNOME'}, _LinuxDesktopEnvironment.GNOME),
  63             ({'XDG_CURRENT_DESKTOP': 'GNOME:GNOME-Classic'}, _LinuxDesktopEnvironment.GNOME),
  64             ({'XDG_CURRENT_DESKTOP': 'GNOME : GNOME-Classic'}, _LinuxDesktopEnvironment.GNOME),
  65
  66             ({'XDG_CURRENT_DESKTOP': 'Unity', 'DESKTOP_SESSION': 'gnome-fallback'}, _LinuxDesktopEnvironment.GNOME),
  67             ({'XDG_CURRENT_DESKTOP': 'KDE', 'KDE_SESSION_VERSION': '5'}, _LinuxDesktopEnvironment.KDE),
  68             ({'XDG_CURRENT_DESKTOP': 'KDE'}, _LinuxDesktopEnvironment.KDE),
  69             ({'XDG_CURRENT_DESKTOP': 'Pantheon'}, _LinuxDesktopEnvironment.PANTHEON),
  70             ({'XDG_CURRENT_DESKTOP': 'Unity'}, _LinuxDesktopEnvironment.UNITY),
  71             ({'XDG_CURRENT_DESKTOP': 'Unity:Unity7'}, _LinuxDesktopEnvironment.UNITY),
  72             ({'XDG_CURRENT_DESKTOP': 'Unity:Unity8'}, _LinuxDesktopEnvironment.UNITY),
  73         ]
  74
  75         for env, expected_desktop_environment in test_cases:
  76             self.assertEqual(_get_linux_desktop_environment(env), expected_desktop_environment)
  77
  78     def test_chrome_cookie_decryptor_linux_derive_key(self):
  79         key = LinuxChromeCookieDecryptor.derive_key(b'abc')
  80         self.assertEqual(key, b'7\xa1\xec\xd4m\xfcA\xc7\xb19Z\xd0\x19\xdcM\x17')
  81
  82     def test_chrome_cookie_decryptor_mac_derive_key(self):
  83         key = MacChromeCookieDecryptor.derive_key(b'abc')
  84         self.assertEqual(key, b'Y\xe2\xc0\xd0P\xf6\xf4\xe1l\xc1\x8cQ\xcb|\xcdY')
  85
  86     def test_chrome_cookie_decryptor_linux_v10(self):
  87         with MonkeyPatch(cookies, {'_get_linux_keyring_password': lambda *args, **kwargs: b''}):
  88             encrypted_value = b'v10\xccW%\xcd\xe6\xe6\x9fM" \xa7\xb0\xca\xe4\x07\xd6'
  89             value = 'USD'
  90             decryptor = LinuxChromeCookieDecryptor('Chrome', Logger())
  91             self.assertEqual(decryptor.decrypt(encrypted_value), value)
  92
  93     def test_chrome_cookie_decryptor_linux_v11(self):
  94         with MonkeyPatch(cookies, {'_get_linux_keyring_password': lambda *args, **kwargs: b''}):
  95             encrypted_value = b'v11#\x81\x10>`w\x8f)\xc0\xb2\xc1\r\xf4\x1al\xdd\x93\xfd\xf8\xf8N\xf2\xa9\x83\xf1\xe9o\x0elVQd'
  96             value = 'tz=Europe.London'
  97             decryptor = LinuxChromeCookieDecryptor('Chrome', Logger())
  98             self.assertEqual(decryptor.decrypt(encrypted_value), value)
  99
 100     def test_chrome_cookie_decryptor_windows_v10(self):
 101         with MonkeyPatch(cookies, {
 102             '_get_windows_v10_key': lambda *args, **kwargs: b'Y\xef\xad\xad\xeerp\xf0Y\xe6\x9b\x12\xc2<z\x16]\n\xbb\xb8\xcb\xd7\x9bA\xc3\x14e\x99{\xd6\xf4&'
 103         }):
 104             encrypted_value = b'v10T\xb8\xf3\xb8\x01\xa7TtcV\xfc\x88\xb8\xb8\xef\x05\xb5\xfd\x18\xc90\x009\xab\xb1\x893\x85)\x87\xe1\xa9-\xa3\xad='
 105             value = '32101439'
 106             decryptor = WindowsChromeCookieDecryptor('', Logger())
 107             self.assertEqual(decryptor.decrypt(encrypted_value), value)
 108
 109     def test_chrome_cookie_decryptor_mac_v10(self):
 110         with MonkeyPatch(cookies, {'_get_mac_keyring_password': lambda *args, **kwargs: b'6eIDUdtKAacvlHwBVwvg/Q=='}):
 111             encrypted_value = b'v10\xb3\xbe\xad\xa1[\x9fC\xa1\x98\xe0\x9a\x01\xd9\xcf\xbfc'
 112             value = '2021-06-01-22'
 113             decryptor = MacChromeCookieDecryptor('', Logger())
 114             self.assertEqual(decryptor.decrypt(encrypted_value), value)
 115
 116     def test_safari_cookie_parsing(self):
 117         cookies = \
 118             b'cook\x00\x00\x00\x01\x00\x00\x00i\x00\x00\x01\x00\x01\x00\x00\x00\x10\x00\x00\x00\x00\x00\x00\x00Y' \
 119             b'\x00\x00\x00\x00\x00\x00\x00 \x00\x00\x00\x00\x00\x00\x008\x00\x00\x00B\x00\x00\x00F\x00\x00\x00H' \
 120             b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x80\x03\xa5>\xc3A\x00\x00\x80\xc3\x07:\xc3A' \
 121             b'localhost\x00foo\x00/\x00test%20%3Bcookie\x00\x00\x00\x054\x07\x17 \x05\x00\x00\x00Kbplist00\xd1\x01' \
 122             b'\x02_\x10\x18NSHTTPCookieAcceptPolicy\x10\x02\x08\x0b&\x00\x00\x00\x00\x00\x00\x01\x01\x00\x00\x00' \
 123             b'\x00\x00\x00\x00\x03\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00('
 124
 125         jar = parse_safari_cookies(cookies)
 126         self.assertEqual(len(jar), 1)
 127         cookie = list(jar)[0]
 128         self.assertEqual(cookie.domain, 'localhost')
 129         self.assertEqual(cookie.port, None)
 130         self.assertEqual(cookie.path, '/')
 131         self.assertEqual(cookie.name, 'foo')
 132         self.assertEqual(cookie.value, 'test%20%3Bcookie')
 133         self.assertFalse(cookie.secure)
 134         expected_expiration = datetime(2021, 6, 18, 21, 39, 19, tzinfo=timezone.utc)
 135         self.assertEqual(cookie.expires, int(expected_expiration.timestamp()))
 136
 137     def test_pbkdf2_sha1(self):
 138         key = pbkdf2_sha1(b'peanuts', b' ' * 16, 1, 16)
 139         self.assertEqual(key, b'g\xe1\x8e\x0fQ\x1c\x9b\xf3\xc9`!\xaa\x90\xd9\xd34')