Add response header in WebSocketBadStatusException

2 files changed, 3 insertions, 2 deletions

diff --git a/websocket/_exceptions.py b/websocket/_exceptions.py
index ee08880..b7a61d3 100644
--- a/websocket/_exceptions.py
+++ b/websocket/_exceptions.py
@@ -74,10 +74,11 @@ class WebSocketBadStatusException(WebSocketException):
     WebSocketBadStatusException will be raised when we get bad handshake status code.
     """
 
-    def __init__(self, message, status_code, status_message=None):
+    def __init__(self, message, status_code, status_message=None, resp_headers=None):
         msg = message % (status_code, status_message)
         super(WebSocketBadStatusException, self).__init__(msg)
         self.status_code = status_code
+        self.resp_headers = resp_headers
 
 class WebSocketAddressException(WebSocketException):
     """
diff --git a/websocket/_handshake.py b/websocket/_handshake.py
index 3b8ddc0..8740fef 100644
--- a/websocket/_handshake.py
+++ b/websocket/_handshake.py
@@ -149,7 +149,7 @@ def _get_handshake_headers(resource, host, port, options):
 def _get_resp_headers(sock, success_statuses=(101, 301, 302, 303)):
     status, resp_headers, status_message = read_headers(sock)
     if status not in success_statuses:
-        raise WebSocketBadStatusException("Handshake status %d %s", status, status_message)
+        raise WebSocketBadStatusException("Handshake status %d %s", status, status_message, resp_headers)
     return status, resp_headers
 
 _HEADERS_TO_CHECK = {