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     LenientSimpleCookie,
   7     LinuxChromeCookieDecryptor,
   8     MacChromeCookieDecryptor,
   9     WindowsChromeCookieDecryptor,
  10     _get_linux_desktop_environment,
  11     _LinuxDesktopEnvironment,
  12     parse_safari_cookies,
  13     pbkdf2_sha1,
  14 )
  15
  16
  17 class Logger:
  18     def debug(self, message, *args, **kwargs):
  19         print(f'[verbose] {message}')
  20
  21     def info(self, message, *args, **kwargs):
  22         print(message)
  23
  24     def warning(self, message, *args, **kwargs):
  25         self.error(message)
  26
  27     def error(self, message, *args, **kwargs):
  28         raise Exception(message)
  29
  30
  31 class MonkeyPatch:
  32     def __init__(self, module, temporary_values):
  33         self._module = module
  34         self._temporary_values = temporary_values
  35         self._backup_values = {}
  36
  37     def __enter__(self):
  38         for name, temp_value in self._temporary_values.items():
  39             self._backup_values[name] = getattr(self._module, name)
  40             setattr(self._module, name, temp_value)
  41
  42     def __exit__(self, exc_type, exc_val, exc_tb):
  43         for name, backup_value in self._backup_values.items():
  44             setattr(self._module, name, backup_value)
  45
  46
  47 class TestCookies(unittest.TestCase):
  48     def test_get_desktop_environment(self):
  49         """ based on https://chromium.googlesource.com/chromium/src/+/refs/heads/main/base/nix/xdg_util_unittest.cc """
  50         test_cases = [
  51             ({}, _LinuxDesktopEnvironment.OTHER),
  52
  53             ({'DESKTOP_SESSION': 'gnome'}, _LinuxDesktopEnvironment.GNOME),
  54             ({'DESKTOP_SESSION': 'mate'}, _LinuxDesktopEnvironment.GNOME),
  55             ({'DESKTOP_SESSION': 'kde4'}, _LinuxDesktopEnvironment.KDE),
  56             ({'DESKTOP_SESSION': 'kde'}, _LinuxDesktopEnvironment.KDE),
  57             ({'DESKTOP_SESSION': 'xfce'}, _LinuxDesktopEnvironment.XFCE),
  58
  59             ({'GNOME_DESKTOP_SESSION_ID': 1}, _LinuxDesktopEnvironment.GNOME),
  60             ({'KDE_FULL_SESSION': 1}, _LinuxDesktopEnvironment.KDE),
  61
  62             ({'XDG_CURRENT_DESKTOP': 'X-Cinnamon'}, _LinuxDesktopEnvironment.CINNAMON),
  63             ({'XDG_CURRENT_DESKTOP': 'GNOME'}, _LinuxDesktopEnvironment.GNOME),
  64             ({'XDG_CURRENT_DESKTOP': 'GNOME:GNOME-Classic'}, _LinuxDesktopEnvironment.GNOME),
  65             ({'XDG_CURRENT_DESKTOP': 'GNOME : GNOME-Classic'}, _LinuxDesktopEnvironment.GNOME),
  66
  67             ({'XDG_CURRENT_DESKTOP': 'Unity', 'DESKTOP_SESSION': 'gnome-fallback'}, _LinuxDesktopEnvironment.GNOME),
  68             ({'XDG_CURRENT_DESKTOP': 'KDE', 'KDE_SESSION_VERSION': '5'}, _LinuxDesktopEnvironment.KDE),
  69             ({'XDG_CURRENT_DESKTOP': 'KDE'}, _LinuxDesktopEnvironment.KDE),
  70             ({'XDG_CURRENT_DESKTOP': 'Pantheon'}, _LinuxDesktopEnvironment.PANTHEON),
  71             ({'XDG_CURRENT_DESKTOP': 'Unity'}, _LinuxDesktopEnvironment.UNITY),
  72             ({'XDG_CURRENT_DESKTOP': 'Unity:Unity7'}, _LinuxDesktopEnvironment.UNITY),
  73             ({'XDG_CURRENT_DESKTOP': 'Unity:Unity8'}, _LinuxDesktopEnvironment.UNITY),
  74         ]
  75
  76         for env, expected_desktop_environment in test_cases:
  77             self.assertEqual(_get_linux_desktop_environment(env), expected_desktop_environment)
  78
  79     def test_chrome_cookie_decryptor_linux_derive_key(self):
  80         key = LinuxChromeCookieDecryptor.derive_key(b'abc')
  81         self.assertEqual(key, b'7\xa1\xec\xd4m\xfcA\xc7\xb19Z\xd0\x19\xdcM\x17')
  82
  83     def test_chrome_cookie_decryptor_mac_derive_key(self):
  84         key = MacChromeCookieDecryptor.derive_key(b'abc')
  85         self.assertEqual(key, b'Y\xe2\xc0\xd0P\xf6\xf4\xe1l\xc1\x8cQ\xcb|\xcdY')
  86
  87     def test_chrome_cookie_decryptor_linux_v10(self):
  88         with MonkeyPatch(cookies, {'_get_linux_keyring_password': lambda *args, **kwargs: b''}):
  89             encrypted_value = b'v10\xccW%\xcd\xe6\xe6\x9fM" \xa7\xb0\xca\xe4\x07\xd6'
  90             value = 'USD'
  91             decryptor = LinuxChromeCookieDecryptor('Chrome', Logger())
  92             self.assertEqual(decryptor.decrypt(encrypted_value), value)
  93
  94     def test_chrome_cookie_decryptor_linux_v11(self):
  95         with MonkeyPatch(cookies, {'_get_linux_keyring_password': lambda *args, **kwargs: b''}):
  96             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'
  97             value = 'tz=Europe.London'
  98             decryptor = LinuxChromeCookieDecryptor('Chrome', Logger())
  99             self.assertEqual(decryptor.decrypt(encrypted_value), value)
 100
 101     def test_chrome_cookie_decryptor_windows_v10(self):
 102         with MonkeyPatch(cookies, {
 103             '_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&'
 104         }):
 105             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='
 106             value = '32101439'
 107             decryptor = WindowsChromeCookieDecryptor('', Logger())
 108             self.assertEqual(decryptor.decrypt(encrypted_value), value)
 109
 110     def test_chrome_cookie_decryptor_mac_v10(self):
 111         with MonkeyPatch(cookies, {'_get_mac_keyring_password': lambda *args, **kwargs: b'6eIDUdtKAacvlHwBVwvg/Q=='}):
 112             encrypted_value = b'v10\xb3\xbe\xad\xa1[\x9fC\xa1\x98\xe0\x9a\x01\xd9\xcf\xbfc'
 113             value = '2021-06-01-22'
 114             decryptor = MacChromeCookieDecryptor('', Logger())
 115             self.assertEqual(decryptor.decrypt(encrypted_value), value)
 116
 117     def test_safari_cookie_parsing(self):
 118         cookies = \
 119             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' \
 120             b'\x00\x00\x00\x00\x00\x00\x00 \x00\x00\x00\x00\x00\x00\x008\x00\x00\x00B\x00\x00\x00F\x00\x00\x00H' \
 121             b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x80\x03\xa5>\xc3A\x00\x00\x80\xc3\x07:\xc3A' \
 122             b'localhost\x00foo\x00/\x00test%20%3Bcookie\x00\x00\x00\x054\x07\x17 \x05\x00\x00\x00Kbplist00\xd1\x01' \
 123             b'\x02_\x10\x18NSHTTPCookieAcceptPolicy\x10\x02\x08\x0b&\x00\x00\x00\x00\x00\x00\x01\x01\x00\x00\x00' \
 124             b'\x00\x00\x00\x00\x03\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00('
 125
 126         jar = parse_safari_cookies(cookies)
 127         self.assertEqual(len(jar), 1)
 128         cookie = list(jar)[0]
 129         self.assertEqual(cookie.domain, 'localhost')
 130         self.assertEqual(cookie.port, None)
 131         self.assertEqual(cookie.path, '/')
 132         self.assertEqual(cookie.name, 'foo')
 133         self.assertEqual(cookie.value, 'test%20%3Bcookie')
 134         self.assertFalse(cookie.secure)
 135         expected_expiration = datetime(2021, 6, 18, 21, 39, 19, tzinfo=timezone.utc)
 136         self.assertEqual(cookie.expires, int(expected_expiration.timestamp()))
 137
 138     def test_pbkdf2_sha1(self):
 139         key = pbkdf2_sha1(b'peanuts', b' ' * 16, 1, 16)
 140         self.assertEqual(key, b'g\xe1\x8e\x0fQ\x1c\x9b\xf3\xc9`!\xaa\x90\xd9\xd34')
 141
 142
 143 class TestLenientSimpleCookie(unittest.TestCase):
 144     def _run_tests(self, *cases):
 145         for message, raw_cookie, expected in cases:
 146             cookie = LenientSimpleCookie(raw_cookie)
 147
 148             with self.subTest(message, expected=expected):
 149                 self.assertEqual(cookie.keys(), expected.keys(), message)
 150
 151                 for key, expected_value in expected.items():
 152                     morsel = cookie[key]
 153                     if isinstance(expected_value, tuple):
 154                         expected_value, expected_attributes = expected_value
 155                     else:
 156                         expected_attributes = {}
 157
 158                     attributes = {
 159                         key: value
 160                         for key, value in dict(morsel).items()
 161                         if value != ""
 162                     }
 163                     self.assertEqual(attributes, expected_attributes, message)
 164
 165                     self.assertEqual(morsel.value, expected_value, message)
 166
 167     def test_parsing(self):
 168         self._run_tests(
 169             # Copied from https://github.com/python/cpython/blob/v3.10.7/Lib/test/test_http_cookies.py
 170             (
 171                 "Test basic cookie",
 172                 "chips=ahoy; vienna=finger",
 173                 {"chips": "ahoy", "vienna": "finger"},
 174             ),
 175             (
 176                 "Test quoted cookie",
 177                 'keebler="E=mc2; L=\\"Loves\\"; fudge=\\012;"',
 178                 {"keebler": 'E=mc2; L="Loves"; fudge=\012;'},
 179             ),
 180             (
 181                 "Allow '=' in an unquoted value",
 182                 "keebler=E=mc2",
 183                 {"keebler": "E=mc2"},
 184             ),
 185             (
 186                 "Allow cookies with ':' in their name",
 187                 "key:term=value:term",
 188                 {"key:term": "value:term"},
 189             ),
 190             (
 191                 "Allow '[' and ']' in cookie values",
 192                 "a=b; c=[; d=r; f=h",
 193                 {"a": "b", "c": "[", "d": "r", "f": "h"},
 194             ),
 195             (
 196                 "Test basic cookie attributes",
 197                 'Customer="WILE_E_COYOTE"; Version=1; Path=/acme',
 198                 {"Customer": ("WILE_E_COYOTE", {"version": "1", "path": "/acme"})},
 199             ),
 200             (
 201                 "Test flag only cookie attributes",
 202                 'Customer="WILE_E_COYOTE"; HttpOnly; Secure',
 203                 {"Customer": ("WILE_E_COYOTE", {"httponly": True, "secure": True})},
 204             ),
 205             (
 206                 "Test flag only attribute with values",
 207                 "eggs=scrambled; httponly=foo; secure=bar; Path=/bacon",
 208                 {"eggs": ("scrambled", {"httponly": "foo", "secure": "bar", "path": "/bacon"})},
 209             ),
 210             (
 211                 "Test special case for 'expires' attribute, 4 digit year",
 212                 'Customer="W"; expires=Wed, 01 Jan 2010 00:00:00 GMT',
 213                 {"Customer": ("W", {"expires": "Wed, 01 Jan 2010 00:00:00 GMT"})},
 214             ),
 215             (
 216                 "Test special case for 'expires' attribute, 2 digit year",
 217                 'Customer="W"; expires=Wed, 01 Jan 98 00:00:00 GMT',
 218                 {"Customer": ("W", {"expires": "Wed, 01 Jan 98 00:00:00 GMT"})},
 219             ),
 220             (
 221                 "Test extra spaces in keys and values",
 222                 "eggs  =  scrambled  ;  secure  ;  path  =  bar   ; foo=foo   ",
 223                 {"eggs": ("scrambled", {"secure": True, "path": "bar"}), "foo": "foo"},
 224             ),
 225             (
 226                 "Test quoted attributes",
 227                 'Customer="WILE_E_COYOTE"; Version="1"; Path="/acme"',
 228                 {"Customer": ("WILE_E_COYOTE", {"version": "1", "path": "/acme"})}
 229             ),
 230             # Our own tests that CPython passes
 231             (
 232                 "Allow ';' in quoted value",
 233                 'chips="a;hoy"; vienna=finger',
 234                 {"chips": "a;hoy", "vienna": "finger"},
 235             ),
 236             (
 237                 "Keep only the last set value",
 238                 "a=c; a=b",
 239                 {"a": "b"},
 240             ),
 241         )
 242
 243     def test_lenient_parsing(self):
 244         self._run_tests(
 245             (
 246                 "Ignore and try to skip invalid cookies",
 247                 'chips={"ahoy;": 1}; vienna="finger;"',
 248                 {"vienna": "finger;"},
 249             ),
 250             (
 251                 "Ignore cookies without a name",
 252                 "a=b; unnamed; c=d",
 253                 {"a": "b", "c": "d"},
 254             ),
 255             (
 256                 "Ignore '\"' cookie without name",
 257                 'a=b; "; c=d',
 258                 {"a": "b", "c": "d"},
 259             ),
 260             (
 261                 "Skip all space separated values",
 262                 "x a=b c=d x; e=f",
 263                 {"a": "b", "c": "d", "e": "f"},
 264             ),
 265             (
 266                 "Skip all space separated values",
 267                 'x a=b; data={"complex": "json", "with": "key=value"}; x c=d x',
 268                 {"a": "b", "c": "d"},
 269             ),
 270             (
 271                 "Expect quote mending",
 272                 'a=b; invalid="; c=d',
 273                 {"a": "b", "c": "d"},
 274             ),
 275             (
 276                 "Reset morsel after invalid to not capture attributes",
 277                 "a=b; invalid; Version=1; c=d",
 278                 {"a": "b", "c": "d"},
 279             ),
 280             (
 281                 "Reset morsel after invalid to not capture attributes",
 282                 "a=b; $invalid; $Version=1; c=d",
 283                 {"a": "b", "c": "d"},
 284             ),
 285             (
 286                 "Continue after non-flag attribute without value",
 287                 "a=b; path; Version=1; c=d",
 288                 {"a": "b", "c": "d"},
 289             ),
 290             (
 291                 "Allow cookie attributes with `$` prefix",
 292                 'Customer="WILE_E_COYOTE"; $Version=1; $Secure; $Path=/acme',
 293                 {"Customer": ("WILE_E_COYOTE", {"version": "1", "secure": True, "path": "/acme"})},
 294             ),
 295             (
 296                 "Invalid Morsel keys should not result in an error",
 297                 "Key=Value; [Invalid]=Value; Another=Value",
 298                 {"Key": "Value", "Another": "Value"},
 299             ),
 300         )