From 64d4c4ba720afdef7ab91c4a77f2e3a063ff5cc0 Mon Sep 17 00:00:00 2001
From: bashonly <bashonly@protonmail.com>
Date: Mon, 8 Apr 2024 15:15:54 -0500
Subject: [PATCH] [networking] Add `ImpersonateResponse`

Authored by: bashonly
---
 yt_dlp/networking/_curlcffi.py   | 28 ++++++++++++++--------------
 yt_dlp/networking/impersonate.py | 14 +++++++++++++-
 2 files changed, 27 insertions(+), 15 deletions(-)

diff --git a/yt_dlp/networking/_curlcffi.py b/yt_dlp/networking/_curlcffi.py
index 39d1f70fb..c7c510ff8 100644
--- a/yt_dlp/networking/_curlcffi.py
+++ b/yt_dlp/networking/_curlcffi.py
@@ -5,13 +5,7 @@ import math
 import urllib.parse
 
 from ._helper import InstanceStoreMixin, select_proxy
-from .common import (
-    Features,
-    Request,
-    Response,
-    register_preference,
-    register_rh,
-)
+from .common import Features, Request, register_preference, register_rh
 from .exceptions import (
     CertificateVerifyError,
     HTTPError,
@@ -20,7 +14,11 @@ from .exceptions import (
     SSLError,
     TransportError,
 )
-from .impersonate import ImpersonateRequestHandler, ImpersonateTarget
+from .impersonate import (
+    ImpersonateRequestHandler,
+    ImpersonateResponse,
+    ImpersonateTarget,
+)
 from ..dependencies import curl_cffi
 from ..utils import int_or_none
 
@@ -80,15 +78,16 @@ class CurlCFFIResponseReader(io.IOBase):
         super().close()
 
 
-class CurlCFFIResponseAdapter(Response):
+class CurlCFFIResponseAdapter(ImpersonateResponse):
     fp: CurlCFFIResponseReader
 
-    def __init__(self, response: curl_cffi.requests.Response):
+    def __init__(self, response: curl_cffi.requests.Response, impersonate: ImpersonateTarget):
         super().__init__(
             fp=CurlCFFIResponseReader(response),
             headers=response.headers,
             url=response.url,
-            status=response.status_code)
+            status=response.status_code,
+            impersonate=impersonate)
 
     def read(self, amt=None):
         try:
@@ -178,6 +177,8 @@ class CurlCFFIRH(ImpersonateRequestHandler, InstanceStoreMixin):
         session.curl.setopt(CurlOpt.LOW_SPEED_LIMIT, 1)  # 1 byte per second
         session.curl.setopt(CurlOpt.LOW_SPEED_TIME, math.ceil(timeout))
 
+        impersonate_target = self._get_request_target(request)
+
         try:
             curl_response = session.request(
                 method=request.method,
@@ -187,8 +188,7 @@ class CurlCFFIRH(ImpersonateRequestHandler, InstanceStoreMixin):
                 verify=self.verify,
                 max_redirects=5,
                 timeout=timeout,
-                impersonate=self._SUPPORTED_IMPERSONATE_TARGET_MAP.get(
-                    self._get_request_target(request)),
+                impersonate=self._SUPPORTED_IMPERSONATE_TARGET_MAP.get(impersonate_target),
                 interface=self.source_address,
                 stream=True
             )
@@ -208,7 +208,7 @@ class CurlCFFIRH(ImpersonateRequestHandler, InstanceStoreMixin):
             else:
                 raise TransportError(cause=e) from e
 
-        response = CurlCFFIResponseAdapter(curl_response)
+        response = CurlCFFIResponseAdapter(curl_response, impersonate_target)
 
         if not 200 <= response.status < 300:
             raise HTTPError(response, redirect_loop=max_redirects_exceeded)
diff --git a/yt_dlp/networking/impersonate.py b/yt_dlp/networking/impersonate.py
index ca66180c7..1e4487d1f 100644
--- a/yt_dlp/networking/impersonate.py
+++ b/yt_dlp/networking/impersonate.py
@@ -5,7 +5,7 @@ from abc import ABC
 from dataclasses import dataclass
 from typing import Any
 
-from .common import RequestHandler, register_preference
+from .common import RequestHandler, Response, register_preference
 from .exceptions import UnsupportedRequest
 from ..compat.types import NoneType
 from ..utils import classproperty, join_nonempty
@@ -58,6 +58,18 @@ class ImpersonateTarget:
         return cls(**mobj.groupdict())
 
 
+class ImpersonateResponse(Response):
+    """
+    Wrapper class for a Response to an impersonated Request.
+
+    Parameters:
+    @param impersonate: the ImpersonateTarget used in the originating Request.
+    """
+    def __init__(self, *args, impersonate: ImpersonateTarget, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.impersonate = impersonate
+
+
 class ImpersonateRequestHandler(RequestHandler, ABC):
     """
     Base class for request handlers that support browser impersonation.