Merge branch 'improve-authcode-receiver' into dev

rayluo · rayluo · commit b53684c6f421 · 2020-11-13T05:52:29.000-08:00
diff --git a/oauth2cli/__init__.py b/oauth2cli/__init__.py
@@ -1,6 +1,7 @@
-__version__ = "0.3.0"
+__version__ = "0.4.0"
 
 from .oidc import Client
 from .assertion import JwtAssertionCreator
 from .assertion import JwtSigner  # Obsolete. For backward compatibility.
+from .authcode import AuthCodeReceiver
 
diff --git a/oauth2cli/authcode.py b/oauth2cli/authcode.py
@@ -5,10 +5,10 @@
 It optionally opens a browser window to guide a human user to manually login.
 After obtaining an auth code, the web server will automatically shut down.
 """
-
-import argparse
 import webbrowser
 import logging
+import socket
+from string import Template
 
 try:  # Python 3
     from http.server import HTTPServer, BaseHTTPRequestHandler
@@ -18,45 +18,23 @@
     from urlparse import urlparse, parse_qs
     from urllib import urlencode
 
-from .oauth2 import Client
-
 
 logger = logging.getLogger(__name__)
 
-def obtain_auth_code(listen_port, auth_uri=None):
-    """This function will start a web server listening on http://localhost:port
-    and then you need to open a browser on this device and visit your auth_uri.
-    When interaction finishes, this function will return the auth code,
-    and then shut down the local web server.
-
-    :param listen_port:
-        The local web server will listen at http://localhost:<listen_port>
-        Unless the authorization server supports dynamic port,
-        you need to use the same port when you register with your app.
-    :param auth_uri: If provided, this function will try to open a local browser.
-    :return: Hang indefinitely, until it receives and then return the auth code.
-    """
-    exit_hint = "Visit http://localhost:{p}?code=exit to abort".format(p=listen_port)
-    logger.warning(exit_hint)
-    if auth_uri:
-        page = "http://localhost:{p}?{q}".format(p=listen_port, q=urlencode({
-            "text": "Open this link to sign in. You may use incognito window",
-            "link": auth_uri,
-            "exit_hint": exit_hint,
-            }))
-        browse(page)
-    server = HTTPServer(("", int(listen_port)), AuthCodeReceiver)
-    try:
-        server.authcode = None
-        while not server.authcode:
-            # Derived from
-            # https://docs.python.org/2/library/basehttpserver.html#more-examples
-            server.handle_request()
-        return server.authcode
-    finally:
-        server.server_close()
 
-def browse(auth_uri):
+def obtain_auth_code(listen_port, auth_uri=None):  # Historically only used in testing
+    with AuthCodeReceiver(port=listen_port) as receiver:
+        return receiver.get_auth_response(
+            auth_uri=auth_uri,
+            welcome_template="""<html><body>
+                Open this link to <a href='$auth_uri'>Sign In</a>
+                (You may want to use incognito window)
+                <hr><a href='$abort_uri'>Abort</a>
+                </body></html>""",
+            ).get("code")
+
+
+def _browse(auth_uri):
     controller = webbrowser.get()  # Get a default controller
     # Some Linux Distro does not setup default browser properly,
     # so we try to explicitly use some popular browser, if we found any.
@@ -69,23 +47,28 @@ def browse(auth_uri):
     logger.info("Please open a browser on THIS device to visit: %s" % auth_uri)
     controller.open(auth_uri)
 
-class AuthCodeReceiver(BaseHTTPRequestHandler):
+
+def _qs2kv(qs):
+    """Flatten parse_qs()'s single-item lists into the item itself"""
+    return {k: v[0] if isinstance(v, list) and len(v) == 1 else v
+        for k, v in qs.items()}
+
+
+class _AuthCodeHandler(BaseHTTPRequestHandler):
     def do_GET(self):
         # For flexibility, we choose to not check self.path matching redirect_uri
         #assert self.path.startswith('/THE_PATH_REGISTERED_BY_THE_APP')
         qs = parse_qs(urlparse(self.path).query)
-        if qs.get('code'):  # Then store it into the server instance
-            ac = self.server.authcode = qs['code'][0]
-            self._send_full_response('Authcode:\n{}'.format(ac))
-            # NOTE: Don't do self.server.shutdown() here. It'll halt the server.
-        elif qs.get('text') and qs.get('link'):  # Then display a landing page
+        if qs.get('code') or qs.get("error"):  # So, it is an auth response
+            self.server.auth_response = _qs2kv(qs)
+            logger.debug("Got auth response: %s", self.server.auth_response)
+            template = (self.server.success_template
+                if "code" in qs else self.server.error_template)
             self._send_full_response(
-                '<a href={link}>{text}</a><hr/>{exit_hint}'.format(
-                link=qs['link'][0], text=qs['text'][0],
-                exit_hint=qs.get("exit_hint", [''])[0],
-                ))
+                template.safe_substitute(**self.server.auth_response))
+            # NOTE: Don't do self.server.shutdown() here. It'll halt the server.
         else:
-            self._send_full_response("This web service serves your redirect_uri")
+            self._send_full_response(self.server.welcome_page)
 
     def _send_full_response(self, body, is_ok=True):
         self.send_response(200 if is_ok else 400)
@@ -95,17 +78,152 @@ def _send_full_response(self, body, is_ok=True):
         self.wfile.write(body.encode("utf-8"))
 
 
+class _AuthCodeHttpServer(HTTPServer):
+    def handle_timeout(self):
+        # It will be triggered when no request comes in self.timeout seconds.
+        # See https://docs.python.org/3/library/socketserver.html#socketserver.BaseServer.handle_timeout
+        raise RuntimeError("Timeout. No auth response arrived.")  # Terminates this server
+            # We choose to not call self.server_close() here,
+            # because it would cause a socket.error exception in handle_request(),
+            # and likely end up the server being server_close() twice.
+
+
+class _AuthCodeHttpServer6(_AuthCodeHttpServer):
+    address_family = socket.AF_INET6
+
+
+class AuthCodeReceiver(object):
+    # This class has (rather than is) an _AuthCodeHttpServer, so it does not leak API
+    def __init__(self, port=None):
+        """Create a Receiver waiting for incoming auth response.
+
+        :param port:
+            The local web server will listen at http://...:<port>
+            You need to use the same port when you register with your app.
+            If your Identity Provider supports dynamic port, you can use port=0 here.
+            Port 0 means to use an arbitrary unused port, per this official example:
+            https://docs.python.org/2.7/library/socketserver.html#asynchronous-mixins
+        """
+        address = "127.0.0.1"  # Hardcode, for now, Not sure what to expose, yet.
+            # Per RFC 8252 (https://tools.ietf.org/html/rfc8252#section-8.3):
+            #   * Clients should listen on the loopback network interface only.
+            #     (It is not recommended to use "" shortcut to bind all addr.)
+            #   * the use of localhost is NOT RECOMMENDED.
+            #     (Use) the loopback IP literal
+            #     rather than localhost avoids inadvertently listening on network
+            #     interfaces other than the loopback interface.
+            # Note:
+            #   When this server physically listens to a specific IP (as it should),
+            #   you will still be able to specify your redirect_uri using either
+            #   IP (e.g. 127.0.0.1) or localhost, whichever matches your registration.
+        Server = _AuthCodeHttpServer6 if ":" in address else _AuthCodeHttpServer
+            # TODO: But, it would treat "localhost" or "" as IPv4.
+            # If pressed, we might just expose a family parameter to caller.
+        self._server = Server((address, port or 0), _AuthCodeHandler)
+
+    def get_port(self):
+        """The port this server actually listening to"""
+        # https://docs.python.org/2.7/library/socketserver.html#SocketServer.BaseServer.server_address
+        return self._server.server_address[1]
+
+    def get_auth_response(self, auth_uri=None, text=None, timeout=None, state=None,
+            welcome_template=None, success_template=None, error_template=None):
+        """Wait and return the auth response, or None when timeout.
+
+        :param str auth_uri:
+            If provided, this function will try to open a local browser.
+        :param str text:
+            If provided (together with auth_uri),
+            this function will render a landing page with ``text`` in your browser.
+            This can be used to make testing more readable.
+        :param int timeout: In seconds. None means wait indefinitely.
+        :param str state:
+            You may provide the state you used in auth_url,
+            then we will use it to validate incoming response.
+        :param str welcome_template:
+            If provided, your end user will see it instead of the auth_uri.
+            When present, it shall be a plaintext or html template following
+            `Python Template string syntax <https://docs.python.org/3/library/string.html#template-strings>`_,
+            and include some of these placeholders: $auth_uri and $abort_uri.
+        :param str success_template:
+            The page will be displayed when authentication was largely successful.
+            Placeholders can be any of these:
+            https://tools.ietf.org/html/rfc6749#section-5.1
+        :param str error_template:
+            The page will be displayed when authentication encountered error.
+            Placeholders can be any of these:
+            https://tools.ietf.org/html/rfc6749#section-5.2
+        :return:
+            The auth response of the first leg of Auth Code flow,
+            typically {"code": "...", "state": "..."} or {"error": "...", ...}
+            See https://tools.ietf.org/html/rfc6749#section-4.1.2
+            and https://openid.net/specs/openid-connect-core-1_0.html#AuthResponse
+            Returns None when the state was mismatched, or when timeout occurred.
+        """
+        welcome_uri = "http://localhost:{p}".format(p=self.get_port())
+        abort_uri = "{loc}?error=abort".format(loc=welcome_uri)
+        logger.debug("Abort by visit %s", abort_uri)
+        self._server.welcome_page = Template(welcome_template or "").safe_substitute(
+            auth_uri=auth_uri, abort_uri=abort_uri)
+        if auth_uri:
+            _browse(welcome_uri if welcome_template else auth_uri)
+        self._server.success_template = Template(success_template or
+            "Authentication completed. You can close this window now.")
+        self._server.error_template = Template(error_template or
+            "Authentication failed. $error: $error_description. ($error_uri)")
+
+        self._server.timeout = timeout  # Otherwise its handle_timeout() won't work
+        self._server.auth_response = {}  # Shared with _AuthCodeHandler
+        while True:
+            # Derived from
+            # https://docs.python.org/2/library/basehttpserver.html#more-examples
+            self._server.handle_request()
+            if self._server.auth_response:
+                if state and state != self._server.auth_response.get("state"):
+                    logger.debug("State mismatch. Ignoring this noise.")
+                else:
+                    break
+        return self._server.auth_response
+
+    def close(self):
+        """Either call this eventually; or use the entire class as context manager"""
+        self._server.server_close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
+# Note: Manually use or test this module by:
+#       python -m path.to.this.file -h
 if __name__ == '__main__':
+    import argparse, json
+    from .oauth2 import Client
     logging.basicConfig(level=logging.INFO)
     p = parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
         description=__doc__ + "The auth code received will be shown at stdout.")
-    p.add_argument('endpoint',
-        help="The auth endpoint for your app. For example: "
-            "https://login.microsoftonline.com/your_tenant/oauth2/authorize")
+    p.add_argument(
+        '--endpoint', help="The auth endpoint for your app.",
+        default="https://login.microsoftonline.com/common/oauth2/v2.0/authorize")
     p.add_argument('client_id', help="The client_id of your application")
-    p.add_argument('redirect_port', type=int, help="The port in redirect_uri")
+    p.add_argument('--port', type=int, default=0, help="The port in redirect_uri")
+    p.add_argument('--host', default="127.0.0.1", help="The host of redirect_uri")
+    p.add_argument('--scope', default=None, help="The scope list")
     args = parser.parse_args()
-    client = Client(args.client_id, authorization_endpoint=args.endpoint)
-    auth_uri = client.build_auth_request_uri("code")
-    print(obtain_auth_code(args.redirect_port, auth_uri))
+    client = Client({"authorization_endpoint": args.endpoint}, args.client_id)
+    with AuthCodeReceiver(port=args.port) as receiver:
+        auth_uri = client.build_auth_request_uri(
+            "code",
+            scope=args.scope.split() if args.scope else None,
+            redirect_uri="http://{h}:{p}".format(h=args.host, p=receiver.get_port()))
+        print(json.dumps(receiver.get_auth_response(
+            auth_uri=auth_uri,
+            welcome_template=
+                "<a href='$auth_uri'>Sign In</a>, or <a href='$abort_uri'>Abort</a",
+            error_template="Oh no. $error",
+            success_template="Oh yeah. Got $code",
+            timeout=60,
+            ), indent=4))
 
