Repository: aaugustin/websockets (python)

Issue:
Support basic auth in client
It does not seem possible to connect to a server using basic auth?

For example:
```
websockets.connect('wss://myuser:mypass@myserver.com')
```

Pull request description:
Add support for HTTP Basic Auth in the client side.
Fix #373.

Patch:
diff --git a/docs/changelog.rst b/docs/changelog.rst
index a961f31..745174f 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -3,11 +3,24 @@ Changelog
 
 .. currentmodule:: websockets
 
-4.1
+5.0
 ...
 
 *In development*
 
+.. warning::
+
+    **Version 5.0 adds a ``user_info`` field to the return value of
+    :func:`~websockets.parse_uri`, :class:`~websockets.WebSocketURI`.**
+
+    If you're unpacking :class:`~websockets.WebSocketURI` into four variables,
+    adjust your code to account for that fifth field.
+
+Also:
+
+* :func:`~client.connect()` performs HTTP Basic Auth when the URI contains
+  credentials.
+
 * :func:`~server.unix_serve` can be used as an asynchronous context manager on
   Python ≥ 3.5.1.
 
diff --git a/websockets/client.py b/websockets/client.py
index 2f69fcd..6922a70 100644
--- a/websockets/client.py
+++ b/websockets/client.py
@@ -16,7 +16,7 @@ from .headers import (
     build_extension_list, build_protocol_list, parse_extension_list,
     parse_protocol_list
 )
-from .http import USER_AGENT, build_headers, read_response
+from .http import USER_AGENT, basic_auth_header, build_headers, read_response
 from .protocol import WebSocketCommonProtocol
 from .uri import parse_uri
 
@@ -225,6 +225,9 @@ class WebSocketClientProtocol(WebSocketCommonProtocol):
         else:
             set_header('Host', '{}:{}'.format(wsuri.host, wsuri.port))
 
+        if wsuri.user_info:
+            set_header(*basic_auth_header(*wsuri.user_info))
+
         if origin is not None:
             set_header('Origin', origin)
 
diff --git a/websockets/http.py b/websockets/http.py
index 99eef84..3fef6d3 100644
--- a/websockets/http.py
+++ b/websockets/http.py
@@ -8,6 +8,7 @@ from :mod:`websockets.http`.
 """
 
 import asyncio
+import base64
 import http.client
 import re
 import sys
@@ -206,3 +207,15 @@ def build_headers(raw_headers):
     headers = http.client.HTTPMessage()
     headers._headers = raw_headers  # HACK
     return headers
+
+
+def basic_auth_header(username, password):
+    """
+    Build an Authorization header for HTTP Basic Auth.
+
+    """
+    # https://tools.ietf.org/html/rfc7617#section-2
+    assert ':' not in username
+    user_pass = '{}:{}'.format(username, password)
+    basic_credentials = base64.b64encode(user_pass.encode()).decode()
+    return ('Authorization', 'Basic ' + basic_credentials)
diff --git a/websockets/uri.py b/websockets/uri.py
index 84c3f3b..21f757f 100644
--- a/websockets/uri.py
+++ b/websockets/uri.py
@@ -15,13 +15,17 @@ from .exceptions import InvalidURI
 __all__ = ['parse_uri', 'WebSocketURI']
 
 WebSocketURI = collections.namedtuple(
-    'WebSocketURI', ['secure', 'host', 'port', 'resource_name'])
+    'WebSocketURI', ['secure', 'host', 'port', 'resource_name', 'user_info'])
 WebSocketURI.__doc__ = """WebSocket URI.
 
 * ``secure`` is the secure flag
 * ``host`` is the lower-case host
 * ``port`` if the integer port, it's always provided even if it's the default
 * ``resource_name`` is the resource name, that is, the path and optional query
+* ``user_info`` is an ``(username, password)`` tuple when the URI contains
+  `User Information`_, else ``None``.
+
+.. _User Information: https://tools.ietf.org/html/rfc3986#section-3.2.1
 
 """
 
@@ -40,8 +44,6 @@ def parse_uri(uri):
         assert uri.scheme in ['ws', 'wss']
         assert uri.params == ''
         assert uri.fragment == ''
-        assert uri.username is None
-        assert uri.password is None
         assert uri.hostname is not None
     except AssertionError as exc:
         raise InvalidURI("{} isn't a valid URI".format(uri)) from exc
@@ -52,4 +54,7 @@ def parse_uri(uri):
     resource_name = uri.path or '/'
     if uri.query:
         resource_name += '?' + uri.query
-    return WebSocketURI(secure, host, port, resource_name)
+    user_info = None
+    if uri.username or uri.password:
+        user_info = (uri.username, uri.password)
+    return WebSocketURI(secure, host, port, resource_name, user_info)


Test patch:
diff --git a/websockets/test_client_server.py b/websockets/test_client_server.py
index e397558..7ee8b78 100644
--- a/websockets/test_client_server.py
+++ b/websockets/test_client_server.py
@@ -109,13 +109,15 @@ def with_client(*args, **kwds):
     return with_manager(temp_test_client, *args, **kwds)
 
 
-def get_server_uri(server, secure=False, resource_name='/'):
+def get_server_uri(server, secure=False, resource_name='/', user_info=None):
     """
     Return a WebSocket URI for connecting to the given server.
 
     """
     proto = 'wss' if secure else 'ws'
 
+    user_info = ':'.join(user_info) + '@' if user_info else ''
+
     # Pick a random socket in order to test both IPv4 and IPv6 on systems
     # where both are available. Randomizing tests is usually a bad idea. If
     # needed, either use the first socket, or test separately IPv4 and IPv6.
@@ -133,7 +135,7 @@ def get_server_uri(server, secure=False, resource_name='/'):
     else:                                                   # pragma: no cover
         raise ValueError("Expected an IPv6, IPv4, or Unix socket")
 
-    return '{}://{}:{}{}'.format(proto, host, port, resource_name)
+    return '{}://{}{}:{}{}'.format(proto, user_info, host, port, resource_name)
 
 
 class UnauthorizedServerProtocol(WebSocketServerProtocol):
@@ -225,11 +227,12 @@ class ClientServerTests(unittest.TestCase):
         start_server = serve(handler, 'localhost', 0, **kwds)
         self.server = self.loop.run_until_complete(start_server)
 
-    def start_client(self, resource_name='/', **kwds):
+    def start_client(self, resource_name='/', user_info=None, **kwds):
         # Don't enable compression by default in tests.
         kwds.setdefault('compression', None)
         secure = kwds.get('ssl') is not None
-        server_uri = get_server_uri(self.server, secure, resource_name)
+        server_uri = get_server_uri(
+            self.server, secure, resource_name, user_info)
         start_client = connect(server_uri, **kwds)
         self.client = self.loop.run_until_complete(start_client)
 
@@ -372,6 +375,14 @@ class ClientServerTests(unittest.TestCase):
         server_path = self.loop.run_until_complete(self.client.recv())
         self.assertEqual(server_path, '/path')
 
+    @with_server()
+    @with_client('/headers', user_info=('user', 'pass'))
+    def test_protocol_basic_auth(self):
+        self.assertEqual(
+            self.client.request_headers['Authorization'],
+            'Basic dXNlcjpwYXNz',
+        )
+
     @with_server()
     @with_client('/headers')
     def test_protocol_headers(self):
diff --git a/websockets/test_http.py b/websockets/test_http.py
index a6d6129..38f6363 100644
--- a/websockets/test_http.py
+++ b/websockets/test_http.py
@@ -2,7 +2,7 @@ import asyncio
 import unittest
 
 from .http import *
-from .http import build_headers, read_headers
+from .http import basic_auth_header, build_headers, read_headers
 
 
 class HTTPAsyncTests(unittest.TestCase):
@@ -128,3 +128,10 @@ class HTTPSyncTests(unittest.TestCase):
 
         # Ordering is deterministic when getting all values.
         self.assertEqual(headers.get_all('X-Foo'), ['Bar', 'Baz'])
+
+    def test_basic_auth_header(self):
+        # Test vector from RFC 7617.
+        self.assertEqual(
+            basic_auth_header("Aladdin", "open sesame"),
+            ('Authorization', 'Basic QWxhZGRpbjpvcGVuIHNlc2FtZQ=='),
+        )
diff --git a/websockets/test_uri.py b/websockets/test_uri.py
index 1b99280..86e305a 100644
--- a/websockets/test_uri.py
+++ b/websockets/test_uri.py
@@ -5,17 +5,32 @@ from .uri import *
 
 
 VALID_URIS = [
-    ('ws://localhost/', (False, 'localhost', 80, '/')),
-    ('wss://localhost/', (True, 'localhost', 443, '/')),
-    ('ws://localhost/path?query', (False, 'localhost', 80, '/path?query')),
-    ('WS://LOCALHOST/PATH?QUERY', (False, 'localhost', 80, '/PATH?QUERY')),
+    (
+        'ws://localhost/',
+        (False, 'localhost', 80, '/', None),
+    ),
+    (
+        'wss://localhost/',
+        (True, 'localhost', 443, '/', None),
+    ),
+    (
+        'ws://localhost/path?query',
+        (False, 'localhost', 80, '/path?query', None),
+    ),
+    (
+        'WS://LOCALHOST/PATH?QUERY',
+        (False, 'localhost', 80, '/PATH?QUERY', None),
+    ),
+    (
+        'ws://user:pass@localhost/',
+        (False, 'localhost', 80, '/', ('user', 'pass')),
+    ),
 ]
 
 INVALID_URIS = [
     'http://localhost/',
     'https://localhost/',
     'ws://localhost/path#fragment',
-    'ws://user:pass@localhost/',
 ]