From e3da5a3ce9189ac9f79ca38502b9c2da456526fe Mon Sep 17 00:00:00 2001 From: Yifei Kong Date: Wed, 26 Jun 2024 11:45:52 +0800 Subject: [PATCH] Add IP info to response --- curl_cffi/requests/models.py | 4 ++++ curl_cffi/requests/session.py | 4 ++-- tests/unittest/test_requests.py | 9 +++++++++ 3 files changed, 15 insertions(+), 2 deletions(-) diff --git a/curl_cffi/requests/models.py b/curl_cffi/requests/models.py index ba94a307..e1af07b4 100644 --- a/curl_cffi/requests/models.py +++ b/curl_cffi/requests/models.py @@ -44,6 +44,8 @@ class Response: elapsed: how many seconds the request cost. encoding: http body encoding. charset: alias for encoding. + primary_ip: primary ip of the server. + local_ip: local ip used in this connection. charset_encoding: encoding specified by the Content-Type header. default_encoding: encoding for decoding response content if charset is not found in headers. Defaults to "utf-8". Can be set to a callable for automatic detection. @@ -68,6 +70,8 @@ def __init__(self, curl: Optional[Curl] = None, request: Optional[Request] = Non self.redirect_count = 0 self.redirect_url = "" self.http_version = 0 + self.primary_ip: str = "" + self.local_ip: str = "" self.history: List[Dict[str, Any]] = [] self.infos: Dict[str, Any] = {} self.queue: Optional[queue.Queue] = None diff --git a/curl_cffi/requests/session.py b/curl_cffi/requests/session.py index 1efe5d50..b029b980 100644 --- a/curl_cffi/requests/session.py +++ b/curl_cffi/requests/session.py @@ -598,11 +598,11 @@ def _parse_response(self, curl, buffer, header_buffer, default_encoding): morsels = [CurlMorsel.from_curl_format(c) for c in c.getinfo(CurlInfo.COOKIELIST)] # for l in c.getinfo(CurlInfo.COOKIELIST): # print("Curl Cookies", l.decode()) - self.cookies.update_cookies_from_curl(morsels) rsp.cookies = self.cookies # print("Cookies after extraction", self.cookies) - + rsp.primary_ip = cast(bytes, c.getinfo(CurlInfo.PRIMARY_IP)).decode() + rsp.local_ip = cast(bytes, c.getinfo(CurlInfo.LOCAL_IP)).decode() rsp.default_encoding = default_encoding rsp.elapsed = cast(float, c.getinfo(CurlInfo.TOTAL_TIME)) rsp.redirect_count = cast(int, c.getinfo(CurlInfo.REDIRECT_COUNT)) diff --git a/tests/unittest/test_requests.py b/tests/unittest/test_requests.py index aa4e3c97..e32d0e11 100644 --- a/tests/unittest/test_requests.py +++ b/tests/unittest/test_requests.py @@ -745,3 +745,12 @@ def test_curl_infos(server): r = s.get(str(server.url)) assert r.infos[CurlInfo.PRIMARY_IP] == b"127.0.0.1" # pyright: ignore + + +def test_response_ip(server): + s = requests.Session() + r = s.get(str(server.url)) + + assert r.primary_ip == "127.0.0.1" + assert r.local_ip == "127.0.0.1" +