fix data stream seek bug

.travis.yml

@@ -10,11 +10,11 @@ install:
 - pip install -U pip setuptools wheel
 - pip install .[testing,linting]
 before_script:
-- python -m tests.test_server &
+- python -m tests.test_server &> /dev/null &
 script:
 - flake8 .
 - bandit -r requests_ntlm2
-- py.test --ignore=tests/functional/test_functional.py --ignore=tests/test_server.py --cov requests_ntlm2 --cov-report term-missing tests
+- py.test -xv --ignore=tests/functional/test_functional.py --ignore=tests/test_server.py --cov requests_ntlm2 --cov-report term-missing tests
 after_success:
 - codecov --token=eb6eedbf-63e1-49ab-a2a9-e8c9e806efe0
 deploy:

requests_ntlm2/connection.py

@@ -9,7 +9,7 @@
 from requests.packages.urllib3.connection import VerifiedHTTPSConnection as _VerifiedHTTPSConnection
 from six.moves.http_client import PROXY_AUTHENTICATION_REQUIRED, LineTooLong
 
-from .core import NtlmCompatibility, get_ntlm_credentials
+from .core import NtlmCompatibility, get_ntlm_credentials, noop
 from .dance import HttpNtlmContext
 
 
@@ -214,6 +214,7 @@ def _tunnel(self):
 
 
 try:
+    noop()  # for testing purposes
     import ssl  # noqa
 
     # Make a copy for testing.

requests_ntlm2/core.py

@@ -186,3 +186,7 @@ def fix_target_info(challenge_msg):
         except struct.error:
             return challenge_msg
     return challenge_msg
+
+
+def noop():
+    pass

requests_ntlm2/requests_ntlm2.py

@@ -1,3 +1,5 @@
+import io
+
 from requests.auth import AuthBase
 
 from .core import NtlmCompatibility, get_auth_type_from_header, get_cbt_data, get_ntlm_credentials
@@ -59,7 +61,10 @@ def retry_using_http_ntlm_auth(
         )
         if hasattr(response.request.body, "seek"):
             if content_length > 0:
-                response.request.body.seek(-content_length, 1)
+                try:
+                    response.request.body.seek(-content_length, 1)
+                except io.UnsupportedOperation:
+                    response.request.body.seek(0, 0)
             else:
                 response.request.body.seek(0, 0)
 

setup.py

@@ -1,7 +1,7 @@
 from setuptools import setup
 
 
-version = "6.2.9"
+version = "6.2.10"
 url = "https://github.com/dopstar/requests-ntlm2"
 
 if "a" in version:

tests/unit/test_connection.py

@@ -1,10 +1,13 @@
 import socket
+import sys
+import tempfile
 import unittest
 
 import faker
 import mock
+from six.moves.http_client import LineTooLong
 
-from requests_ntlm2.connection import VerifiedHTTPSConnection
+from requests_ntlm2.connection import _MAXLINE, VerifiedHTTPSConnection
 
 
 try:
@@ -79,6 +82,37 @@ def test__tunnel__no_credentials(self):
         with self.assertRaises(AttributeError):
             self.conn._tunnel()
 
+    @mock.patch("requests_ntlm2.connection.select.select")
+    @mock.patch("requests_ntlm2.connection.VerifiedHTTPSConnection._get_response")
+    @mock.patch("requests_ntlm2.connection.VerifiedHTTPSConnection.send")
+    def test__tunnel__line_too_long(self, mock_send, mock_get_response, mock_select):
+        fp = BytesIO(
+            b"Proxy-Authenticate: NTLM TlRMTVNTUAACAAAABgAGADgAAAAGgokAyYpGWqVMA/QAAAAAAAAA"
+            b"AH4AfgA+AAAABQCTCAAAAA9ERVROU1cCAAwARABFAFQATgBTAFcAAQAaAFMARwAtADQAOQAxADMAM"
+            b"wAwADAAMAAwADkABAAUAEQARQBUAE4AUwBXAC4AVwBJAE4AAwAwAHMAZwAtADQAOQAxADMAMwAwAD"
+            b"AAMAAwADkALgBkAGUAdABuAHMAdwAuAHcAaQBuAAAAAAA=\r\n"
+            b"Connection: Keep-Alive\r\n"
+            b"Proxy-Connection: Keep-Alive\r\n"
+            b"Server: nginx\r\n"
+            b"Info: x%s\r\n"
+            b"\r\n"
+            b"this is the body\r\n"
+            b"\r\n" % (b"x" * _MAXLINE)
+        )
+        response = type("Response", (), dict(fp=fp))
+        mock_get_response.return_value = "HTTP/1.1", 407, "Proxy Authentication Required", response
+        mock_select.return_value = [(True), (), ()]
+        self.conn.set_ntlm_auth_credentials(r"DOMAIN\username", "password")
+
+        with self.assertRaises(LineTooLong):
+            self.conn._tunnel()
+
+        fp.seek(0)
+        mock_get_response.return_value = "HTTP/1.1", 200, "OK", response
+        self.conn.set_ntlm_auth_credentials(r"DOMAIN\username", "password")
+        with self.assertRaises(LineTooLong):
+            self.conn._tunnel()
+
     @mock.patch("requests_ntlm2.connection.select.select")
     @mock.patch("requests_ntlm2.connection.VerifiedHTTPSConnection._get_response")
     @mock.patch("requests_ntlm2.connection.VerifiedHTTPSConnection.send")
@@ -93,7 +127,7 @@ def test_tunnel__no_headers(self, mock_send, mock_get_response, mock_select):
         self.conn.set_ntlm_auth_credentials(username, password)
 
         error_msg = "Tunnel connection failed: 407 Proxy Authentication Required"
-        with self.assertRaisesRegexp(socket.error, error_msg):
+        with self.assertRaisesRegex(socket.error, error_msg):
             self.conn._tunnel()
 
         mock_get_response.assert_called()
@@ -296,3 +330,110 @@ def test_handle_http09__worst_case(self):
         response = type("Response", (), dict(fp=fp))
         status_line = self.conn.handle_http09_response(response)
         self.assertIsNone(status_line)
+
+    def test__read_response_line_if_ready(self):
+        data = (
+            b"<!DOCTYPE html>\r\n"
+            b"<html class='#{theme}' lang='en'>\r\n"
+            b"<head data-theme='#{theme}' data-revision='865b887'>\r\n"
+            b"<meta charset='utf-8'/>\r\n"
+            b"<meta http-equiv='X-UA-Compatible' content='IE=edge'/>\r\n"
+            b"<meta name='viewport' content='width=device-width, initial-scale=1'/>\r\n"
+            b"<base/>\r\n"
+            b"<title>401 Unauthorised</title><!--[if lt IE 9]>\r\n"
+            b"<script src=\'https://oss.maxcdn.com/html5shiv/3.7.2/html5shiv.min.js\' type=\'text/javascript\' />\r\n"  # noqa
+            b"<script src=\'https://oss.maxcdn.com/respond/1.4.2/respond.min.js\' type=\'text/javascript\' />\r\n"  # noqa
+            b"<![endif]-->\r\n"
+            b"<script type='text/javascript'>\r\n"
+            b"function showURL()\r\n"
+            b"{\r\n"
+            b"document.write('<a href=\"mailto:[email protected]?subject=Proxy Authentication\">EDConnect</a>');\r\n"  # noqa
+            b"}\r\n"
+            b"\r\n"
+            b"function URL()\r\n"
+            b"{\r\n"
+            b"document.write(document.URL);\r\n"
+            b"}\r\n"
+            b"</script>\r\n"
+            b"<style media='screen'>\r\n"
+            b"</div>\r\n"
+            b"</div>\r\n"
+            b"</div>\r\n"
+            b"</aside></body>\r\n"
+            b"</html>\r\n"
+        )
+        with tempfile.TemporaryFile() as fd:
+            fd.write(data)
+            fd.seek(0)
+            response = type("Response", (), dict(fp=fd))
+            line = self.conn._read_response_line_if_ready(response)
+            assert line == b"<!DOCTYPE html>\r\n"
+
+        with tempfile.TemporaryFile() as fd:
+            fd.write(data)
+            fd.seek(0)
+            response = type("Response", (), dict(fp=fd))
+            with mock.patch("select.select", return_value=((), (), ())) as mock_select:
+                line = self.conn._read_response_line_if_ready(response)
+                assert line is None
+                mock_select.assert_called_once_with([response.fp], (), (), 0.1)
+
+    def test__flush_response_buffer(self):
+        data = (
+            b"<!DOCTYPE html>\r\n"
+            b"<html class='#{theme}' lang='en'>\r\n"
+            b"<head data-theme='#{theme}' data-revision='865b887'>\r\n"
+            b"<meta charset='utf-8'/>\r\n"
+            b"<meta http-equiv='X-UA-Compatible' content='IE=edge'/>\r\n"
+            b"<meta name='viewport' content='width=device-width, initial-scale=1'/>\r\n"
+            b"<base/>\r\n"
+            b"<title>401 Unauthorised</title><!--[if lt IE 9]>\r\n"
+            b"<script src=\'https://oss.maxcdn.com/html5shiv/3.7.2/html5shiv.min.js\' type=\'text/javascript\' />\r\n"  # noqa
+            b"<script src=\'https://oss.maxcdn.com/respond/1.4.2/respond.min.js\' type=\'text/javascript\' />\r\n"  # noqa
+            b"<![endif]-->\r\n"
+            b"<script type='text/javascript'>\r\n"
+            b"function showURL()\r\n"
+            b"{\r\n"
+            b"document.write('<a href=\"mailto:[email protected]?subject=Proxy Authentication\">EDConnect</a>');\r\n"  # noqa
+            b"}\r\n"
+            b"\r\n"
+            b"function URL()\r\n"
+            b"{\r\n"
+            b"document.write(document.URL);\r\n"
+            b"}\r\n"
+            b"</script>\r\n"
+            b"<style media='screen'>\r\n"
+            b"</div>\r\n"
+            b"</div>\r\n"
+            b"</div>\r\n"
+            b"</aside></body>\r\n"
+            b"</html>\r\n"
+        )
+        with tempfile.TemporaryFile() as fd:
+            fd.write(data)
+            fd.seek(0)
+            response = type("Response", (), dict(fp=fd))
+            result = self.conn._flush_response_buffer(response)
+            assert result is None
+            assert fd.read() == b""
+
+        with tempfile.TemporaryFile() as fd:
+            fd.write(data)
+            fd.seek(0)
+            response = type("Response", (), dict(fp=fd))
+            with mock.patch("select.select", return_value=((), (), ())) as mock_select:
+                result = self.conn._flush_response_buffer(response)
+                assert result is None
+                mock_select.assert_called_once_with([response.fp], (), (), 0.1)
+                assert fd.read() == data
+
+
+def test_import_error():
+    with mock.patch("requests_ntlm2.core.noop") as mock_noop:
+        mock_noop.side_effect = ImportError()
+        import requests_ntlm2.connection  # noqa - this is ensure sys.modules key is present
+        del sys.modules["requests_ntlm2.connection"]
+        import requests_ntlm2.connection  # noqa
+        assert (
+            requests_ntlm2.connection.HTTPSConnection is requests_ntlm2.connection.DummyConnection
+        )

tests/unit/test_dance.py

@@ -15,7 +15,7 @@ def setUp(self):
 
     def test__init(self):
         error_msg = 'Expected "NTLM" or "Negotiate" auth_type, got None'
-        with self.assertRaisesRegexp(ValueError, error_msg):
+        with self.assertRaisesRegex(ValueError, error_msg):
             requests_ntlm2.dance.HttpNtlmContext("username", "password")
 
     def test__init__ntlm(self):