Add Okta Login Handler (#1)

* Add Okta Login Handler

docs/auth.rst

@@ -53,6 +53,27 @@ NOTE: Enable Google Plus API in the Google Developers Console under `APIs & auth
 
 .. _github-oauth:
 
+Okta OAuth
+------------
+
+Flower also supports Okta OAuth. Flower should be registered in
+<https://developer.okta.com/docs/guides/add-an-external-idp/openidconnect/register-app-in-okta/> before getting started.
+See `Okta OAuth API`_ docs for more info.
+
+Okta OAuth should be activated using `--auth_provider` option.
+The client id, secret and redirect uri should be provided using
+`--oauth2_key`, `--oauth2_secret`, `--oauth2_redirect_uri`, `--oauth2_okta_base_url` options or using
+`FLOWER_OAUTH2_KEY`, `FLOWER_OAUTH2_SECRET`, `FLOWER_OAUTH2_REDIRECT_URI`, and `FLOWER_OAUTH2_OKTA_BASE_URL`
+environment variables. ::
+
+    $ export FLOWER_OAUTH2_KEY=7956724aafbf5e1a93ac
+    $ export FLOWER_OAUTH2_SECRET=f9155f764b7e466c445931a6e3cc7a42c4ce47be
+    $ export FLOWER_OAUTH2_REDIRECT_URI=http://localhost:5555/login
+    $ export FLOWER_OAUTH2_OKTA_BASE_URL=https://my-company.okta.com/oauth2
+    $ celery flower --auth_provider=flower.views.auth.OktaLoginHandler --auth=.*@example\.com
+
+.. _Okta OAuth API: https://developer.okta.com/docs/reference/api/oidc/
+
 GitHub OAuth
 ------------
 

flower/command.py

@@ -106,6 +106,7 @@ def extract_settings(self):
                 'key': options.oauth2_key or os.environ.get('FLOWER_OAUTH2_KEY'),
                 'secret': options.oauth2_secret or os.environ.get('FLOWER_OAUTH2_SECRET'),
                 'redirect_uri': options.oauth2_redirect_uri or os.environ.get('FLOWER_OAUTH2_REDIRECT_URI'),
+                'okta_base_url': options.oauth2_okta_base_url or os.environ.get('FLOWER_OAUTH2_OKTA_BASE_URL'),
             }
 
         if options.certfile and options.keyfile:

flower/options.py

@@ -29,6 +29,9 @@
        help="OAuth2 secret (requires --auth)")
 define("oauth2_redirect_uri", type=str, default=None,
        help="OAuth2 redirect uri (requires --auth)")
+define("oauth2_okta_base_url", type=str, default=None,
+       help="Okta's OAuth2 base url, e.g. https://example.okta.com/oauth2 "
+            "(requires --auth_provider=flower.views.auth.OktaLoginHandler)")
 define("max_workers", type=int, default=5000,
        help="maximum number of workers to keep in memory")
 define("max_tasks", type=int, default=10000,

flower/views/auth.py

@@ -1,7 +1,9 @@
 from __future__ import absolute_import
 
 import json
+import os
 import re
+import uuid
 
 try:
     from urllib.parse import urlencode
@@ -150,6 +152,109 @@ def _on_auth(self, user):
         self.redirect(next_)
 
 
+class OktaLoginHandler(BaseHandler, tornado.auth.OAuth2Mixin):
+
+    _OAUTH_NO_CALLBACKS = False
+    _OAUTH_SETTINGS_KEY = 'oauth'
+
+    @property
+    def base_url(self):
+        return self.settings[self._OAUTH_SETTINGS_KEY]['okta_base_url']
+
+    @property
+    def _OAUTH_AUTHORIZE_URL(self):
+        return "{}/v1/authorize".format(self.base_url)
+
+    @property
+    def _OAUTH_ACCESS_TOKEN_URL(self):
+        return "{}/v1/token".format(self.base_url)
+
+    @property
+    def _OAUTH_USER_INFO_URL(self):
+        return "{}/v1/userinfo".format(self.base_url)
+
+    @tornado.gen.coroutine
+    def get_access_token(self, redirect_uri, code):
+        body = urlencode({
+            "redirect_uri": redirect_uri,
+            "code": code,
+            "client_id": self.settings[self._OAUTH_SETTINGS_KEY]['key'],
+            "client_secret": self.settings[self._OAUTH_SETTINGS_KEY]['secret'],
+            "grant_type": "authorization_code",
+        })
+
+        response = yield self.get_auth_http_client().fetch(
+            self._OAUTH_ACCESS_TOKEN_URL,
+            method="POST",
+            headers={'Content-Type': 'application/x-www-form-urlencoded',
+                     'Accept': 'application/json'}, body=body)
+
+        if response.error:
+            raise tornado.auth.AuthError(
+                'OAuth authenticator error: %s' % str(response))
+
+        raise tornado.gen.Return(json.loads(response.body.decode('utf-8')))
+
+    @tornado.gen.coroutine
+    def get(self):
+        redirect_uri = self.settings[self._OAUTH_SETTINGS_KEY]['redirect_uri']
+        if self.get_argument('code', False):
+            expected_state = self.get_secure_cookie('oauth_state').decode('utf-8')
+            returned_state = self.get_argument('state')
+
+            if returned_state is None or returned_state != expected_state:
+                raise tornado.auth.AuthError(
+                    'OAuth authenticator error: State tokens do not match')
+
+            access_token_response = yield self.get_access_token(
+                redirect_uri=redirect_uri,
+                code=self.get_argument('code'),
+            )
+            yield self._on_auth(access_token_response)
+        else:
+            state = str(uuid.uuid4())
+            self.set_secure_cookie("oauth_state", state)
+            yield self.authorize_redirect(
+                redirect_uri=redirect_uri,
+                client_id=self.settings[self._OAUTH_SETTINGS_KEY]['key'],
+                scope=['openid email'],
+                response_type='code',
+                extra_params={'state': state}
+            )
+
+    @tornado.gen.coroutine
+    def _on_auth(self, access_token_response):
+        if not access_token_response:
+            raise tornado.web.HTTPError(500, 'OAuth authentication failed')
+        access_token = access_token_response['access_token']
+
+        response = yield self.get_auth_http_client().fetch(
+            self._OAUTH_USER_INFO_URL,
+            headers={'Authorization': 'Bearer ' + access_token,
+                     'User-agent': 'Tornado auth'})
+
+        decoded_body = json.loads(response.body.decode('utf-8'))
+        email = (decoded_body.get('email') or '').strip()
+        email_verified = (
+            decoded_body.get('email_verified') and
+            re.match(self.application.options.auth, email)
+        )
+
+        if not email_verified:
+            message = (
+                "Access denied. Please use another account or "
+                "ask your admin to add your email to flower --auth."
+            )
+            raise tornado.web.HTTPError(403, message)
+
+        self.set_secure_cookie("user", str(email))
+
+        next_ = self.get_argument('next', self.application.options.url_prefix or '/')
+        if self.application.options.url_prefix and next_[0] != '/':
+            next_ = '/' + next_
+        self.redirect(next_)
+
+
 class LogoutHandler(BaseHandler):
     def get(self):
         self.clear_cookie('user')