From a99d5b9e015bb4807a4dd9215a1ba137f37dcd85 Mon Sep 17 00:00:00 2001
From: debugtalk <mail@debugtalk.com>
Date: Wed, 17 Jun 2020 21:20:45 +0800
Subject: [PATCH] feat: log client/server IP and port

---
 docs/CHANGELOG.md    |  4 ++++
 httprunner/client.py | 27 +++++++++++++++++++++------
 httprunner/make.py   |  4 ++--
 httprunner/models.py |  8 ++++++++
 4 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md
index 72f6656a..9c36fb26 100644
--- a/docs/CHANGELOG.md
+++ b/docs/CHANGELOG.md
@@ -2,6 +2,10 @@
 
 ## 3.0.13 (2020-06-15)
 
+**Added**
+
+- feat: log client/server IP and port
+
 **Fixed**
 
 - fix: avoid '.csv' been converted to '_csv'
diff --git a/httprunner/client.py b/httprunner/client.py
index 1580fecf..b7cf1447 100644
--- a/httprunner/client.py
+++ b/httprunner/client.py
@@ -172,16 +172,31 @@ class HttpSession(requests.Session):
         # timeout default to 120 seconds
         kwargs.setdefault("timeout", 120)
 
+        # set stream to True, in order to get client/server IP/Port
+        kwargs["stream"] = True
+
         start_timestamp = time.time()
         response = self._send_request_safe_mode(method, url, **kwargs)
         response_time_ms = round((time.time() - start_timestamp) * 1000, 2)
 
-        # get the length of the content, but if the argument stream is set to True, we take
-        # the size from the content-length header, in order to not trigger fetching of the body
-        if kwargs.get("stream", False):
-            content_size = int(dict(response.headers).get("content-length") or 0)
-        else:
-            content_size = len(response.content or "")
+        try:
+            client_ip, client_port = response.raw.connection.sock.getsockname()
+            self.data.address.client_ip = client_ip
+            self.data.address.client_port = client_port
+            logger.debug(f"client IP: {client_ip}, Port: {client_port}")
+        except AttributeError:
+            logger.warning(f"failed to get client address info!")
+
+        try:
+            server_ip, server_port = response.raw.connection.sock.getpeername()
+            self.data.address.server_ip = server_ip
+            self.data.address.server_port = server_port
+            logger.debug(f"server IP: {server_ip}, Port: {server_port}")
+        except AttributeError:
+            logger.warning(f"failed to get server address info!")
+
+        # get length of the response content
+        content_size = int(dict(response.headers).get("content-length") or 0)
 
         # record the consumed time
         self.data.stat.response_time_ms = response_time_ms
diff --git a/httprunner/make.py b/httprunner/make.py
index ef4ce731..c4df5a93 100644
--- a/httprunner/make.py
+++ b/httprunner/make.py
@@ -66,10 +66,10 @@ if __name__ == "__main__":
 def __ensure_absolute(path: Text) -> Text:
     if path.startswith("./"):
         # Linux/Darwin, hrun ./test.yml
-        path = path[len("./"):]
+        path = path[len("./") :]
     elif path.startswith(".\\"):
         # Windows, hrun .\\test.yml
-        path = path[len(".\\"):]
+        path = path[len(".\\") :]
 
     path = ensure_path_sep(path)
     project_meta = load_project_meta(path)
diff --git a/httprunner/models.py b/httprunner/models.py
index 61061bff..0e20a494 100644
--- a/httprunner/models.py
+++ b/httprunner/models.py
@@ -110,6 +110,13 @@ class RequestStat(BaseModel):
     elapsed_ms: float = 0
 
 
+class AddressData(BaseModel):
+    client_ip: Text = "N/A"
+    client_port: int = 0
+    server_ip: Text = "N/A"
+    server_port: int = 0
+
+
 class RequestData(BaseModel):
     method: MethodEnum = MethodEnum.GET
     url: Url
@@ -140,6 +147,7 @@ class SessionData(BaseModel):
     # while when 30X redirect occurs, req_resps will contain multiple request & response
     req_resps: List[ReqRespData] = []
     stat: RequestStat = RequestStat()
+    address: AddressData = AddressData()
     validators: Dict = {}