Merge pull request jupyter#58 from jbweston/feature/group-whitelisting

Group/Organization whitelisting

oauthenticator/bitbucket.py

@@ -46,6 +46,13 @@ class BitbucketOAuthenticator(OAuthenticator):
         help="Automatically whitelist members of selected teams",
     )
 
+    bitbucket_team_whitelist = team_whitelist
+
+
+    headers = {"Accept": "application/json",
+               "User-Agent": "JupyterHub",
+               "Authorization": "Bearer {}"
+               }
 
     @gen.coroutine
     def authenticate(self, handler, data=None):
@@ -95,7 +102,7 @@ def authenticate(self, handler, data=None):
 
         # Check if user is a member of any whitelisted teams.
         # This check is performed here, as the check requires `access_token`.
-        if self.team_whitelist:
+        if self.bitbucket_team_whitelist:
             user_in_team = yield self._check_team_whitelist(username, access_token)
             return username if user_in_team else None
         else:  # no team whitelisting
@@ -109,16 +116,18 @@ def _check_team_whitelist(self, username, access_token):
         # We verify the team membership by calling teams endpoint.
         next_page = url_concat("https://api.bitbucket.org/2.0/teams",
                                {'role': 'member'})
-        user_teams = set()
         while next_page:
             req = HTTPRequest(next_page, method="GET", headers=headers)
             resp = yield http_client.fetch(req)
             resp_json = json.loads(resp.body.decode('utf8', 'replace'))
             next_page = resp_json.get('next', None)
 
-            user_teams |= \
+            user_teams = \
                 set([entry["username"] for entry in resp_json["values"]])
-        return len(self.team_whitelist & user_teams) > 0
+            # check if any of the organizations seen thus far are in whitelist
+            if len(self.bitbucket_team_whitelist & user_teams) > 0:
+                return True
+        return False
 
 
 class LocalBitbucketOAuthenticator(LocalAuthenticator,

oauthenticator/github.py

@@ -7,25 +7,49 @@
 
 import json
 import os
+import re
 
 from tornado.auth import OAuth2Mixin
 from tornado import gen, web
 
+import requests
 from tornado.httputil import url_concat
 from tornado.httpclient import HTTPRequest, AsyncHTTPClient
 
 from jupyterhub.auth import LocalAuthenticator
 
-from traitlets import Unicode
+from traitlets import Unicode, Set
 
 from .oauth2 import OAuthLoginHandler, OAuthenticator
 
 # Support github.com and github enterprise installations
 GITHUB_HOST = os.environ.get('GITHUB_HOST') or 'github.com'
 if GITHUB_HOST == 'github.com':
-    GITHUB_API = 'api.github.com/user'
+    GITHUB_API = 'api.github.com'
 else:
-    GITHUB_API = '%s/api/v3/user' % GITHUB_HOST
+    GITHUB_API = '%s/api/v3' % GITHUB_HOST
+
+
+def _api_headers(access_token):
+    return {"Accept": "application/json",
+            "User-Agent": "JupyterHub",
+            "Authorization": "token {}".format(access_token)
+           }
+
+
+def _get_next_page(response):
+    # Github uses Link headers for pagination.
+    # See https://developer.github.com/v3/#pagination
+    link_header = response.headers.get('Link')
+    if not link_header:
+        return
+    for link in requests.utils.parse_header_links(link_header):
+        if link.get('rel') == 'next':
+            return link['url']
+    # if no "next" page, this is the last one
+    return None
+
+
 
 class GitHubMixin(OAuth2Mixin):
     _OAUTH_AUTHORIZE_URL = "https://%s/login/oauth/authorize" % GITHUB_HOST
@@ -53,6 +77,12 @@ def _github_client_secret_changed(self, name, old, new):
     client_id_env = 'GITHUB_CLIENT_ID'
     client_secret_env = 'GITHUB_CLIENT_SECRET'
     login_handler = GitHubLoginHandler
+
+    github_organization_whitelist = Set(
+        config=True,
+        help="Automatically whitelist members of selected organizations",
+    )
+
 
     @gen.coroutine
     def authenticate(self, handler, data=None):
@@ -88,18 +118,45 @@ def authenticate(self, handler, data=None):
         access_token = resp_json['access_token']
 
         # Determine who the logged in user is
-        headers={"Accept": "application/json",
-                 "User-Agent": "JupyterHub",
-                 "Authorization": "token {}".format(access_token)
-        }
-        req = HTTPRequest("https://%s" % GITHUB_API,
+        req = HTTPRequest("https://%s/user" % GITHUB_API,
                           method="GET",
-                          headers=headers
+                          headers=_api_headers(access_token)
                           )
         resp = yield http_client.fetch(req)
         resp_json = json.loads(resp.body.decode('utf8', 'replace'))
 
-        return resp_json["login"]
+        username =  resp_json["login"]
+
+        # Check if user is a member of any whitelisted organizations.
+        # This check is performed here, as it requires `access_token`.
+        if self.github_organization_whitelist:
+            for org in self.github_organization_whitelist:
+                user_in_org = yield self._check_organization_whitelist(org, username, access_token)
+                if user_in_org:
+                    return username
+            else:  # User not found in member list for any organisation
+                return None
+        else:  # no organization whitelisting
+            return username
+
+
+    @gen.coroutine
+    def _check_organization_whitelist(self, org, username, access_token):
+        http_client = AsyncHTTPClient()
+        headers = _api_headers(access_token)
+        # Get all the members for organization 'org'
+        next_page = "https://%s/orgs/%s/members" % (GITHUB_API, org)
+        while next_page:
+            req = HTTPRequest(next_page, method="GET", headers=headers)
+            resp = yield http_client.fetch(req)
+            resp_json = json.loads(resp.body.decode('utf8', 'replace'))
+            next_page = _get_next_page(resp)
+            org_members = set(entry["login"] for entry in resp_json)
+            # check if any of the organizations seen thus far are in whitelist
+            if username in org_members:
+                return True
+        return False
+
 
 
 class LocalGitHubOAuthenticator(LocalAuthenticator, GitHubOAuthenticator):

oauthenticator/gitlab.py

@@ -13,16 +13,41 @@
 from tornado.auth import OAuth2Mixin
 from tornado import gen, web
 
+import requests
+from tornado.escape import url_escape
 from tornado.httputil import url_concat
 from tornado.httpclient import HTTPRequest, AsyncHTTPClient
 
 from jupyterhub.auth import LocalAuthenticator
 
+from traitlets import Set
+
 from .oauth2 import OAuthLoginHandler, OAuthenticator
 
 # Support gitlab.com and gitlab community edition installations
 GITLAB_HOST = os.environ.get('GITLAB_HOST') or 'https://gitlab.com'
-GITLAB_API = '%s/api/v3/user' % GITLAB_HOST
+GITLAB_API = '%s/api/v3' % GITLAB_HOST
+
+
+def _api_headers(access_token):
+    return {"Accept": "application/json",
+            "User-Agent": "JupyterHub",
+            "Authorization": "token {}".format(access_token)
+           }
+
+
+def _get_next_page(response):
+    # Gitlab uses Link headers for pagination.
+    # See https://docs.gitlab.com/ee/api/README.html#pagination-link-header
+    link_header = response.headers.get('Link')
+    if not link_header:
+        return
+    for link in requests.utils.parse_header_links(link_header):
+        if link.get('rel') == 'next':
+            return link['url']
+    # if no "next" page, this is the last one
+    return None
+
 
 class GitLabMixin(OAuth2Mixin):
     _OAUTH_AUTHORIZE_URL = "%s/oauth/authorize" % GITLAB_HOST
@@ -41,6 +66,12 @@ class GitLabOAuthenticator(OAuthenticator):
     client_secret_env = 'GITLAB_CLIENT_SECRET'
     login_handler = GitLabLoginHandler
 
+    gitlab_group_whitelist = Set(
+        config=True,
+        help="Automatically whitelist members of selected groups",
+    )
+
+
     @gen.coroutine
     def authenticate(self, handler, data=None):
         code = handler.get_argument("code", False)
@@ -81,18 +112,57 @@ def authenticate(self, handler, data=None):
         access_token = resp_json['access_token']
 
         # Determine who the logged in user is
-        headers={"Accept": "application/json",
-                 "User-Agent": "JupyterHub",
-        }
-        req = HTTPRequest("%s?access_token=%s" % (GITLAB_API, access_token),
+        req = HTTPRequest("%s/user" % GITLAB_API,
                           method="GET",
                           validate_cert=validate_server_cert,
-                          headers=headers
+                          headers=_api_headers(access_token)
                           )
         resp = yield http_client.fetch(req)
         resp_json = json.loads(resp.body.decode('utf8', 'replace'))
 
-        return resp_json["username"]
+        username = resp_json["username"]
+        user_id = resp_json["id"]
+        is_admin = resp_json["is_admin"]
+
+        # Check if user is a member of any whitelisted organizations.
+        # This check is performed here, as it requires `access_token`.
+        if self.gitlab_group_whitelist:
+            user_in_group = yield self._check_group_whitelist(
+                username, user_id, is_admin, access_token)
+            return username if user_in_group else None
+        else:  # no organization whitelisting
+            return username
+
+
+    @gen.coroutine
+    def _check_group_whitelist(self, username, user_id, is_admin, access_token):
+        http_client = AsyncHTTPClient()
+        headers = _api_headers(access_token)
+        if is_admin:
+            # For admins, /groups returns *all* groups. As a workaround
+            # we check if we are a member of each group in the whitelist
+            for group in map(url_escape, self.gitlab_group_whitelist):
+                url = "%s/groups/%s/members/%d" % (GITLAB_API, group, user_id)
+                req = HTTPRequest(url, method="GET", headers=headers)
+                resp = yield http_client.fetch(req, raise_error=False)
+                if resp.code == 200:
+                    return True  # user _is_ in group
+        else:
+            # For regular users we get all the groups to which they have access
+            # and check if any of these are in the whitelisted groups
+            next_page = url_concat("%s/groups" % GITLAB_API,
+                                   dict(all_available=True))
+            while next_page:
+                req = HTTPRequest(next_page, method="GET", headers=headers)
+                resp = yield http_client.fetch(req)
+                resp_json = json.loads(resp.body.decode('utf8', 'replace'))
+                next_page = _get_next_page(resp)
+                user_groups = set(entry["path"] for entry in resp_json)
+                # check if any of the organizations seen thus far are in whitelist
+                if len(self.gitlab_group_whitelist & user_groups) > 0:
+                    return True
+            return False
+
 
 
 class LocalGitLabOAuthenticator(LocalAuthenticator, GitLabOAuthenticator):

oauthenticator/tests/test_bitbucket.py

@@ -41,7 +41,7 @@ def test_no_code(bitbucket_client):
 def test_team_whitelist(bitbucket_client):
     client = bitbucket_client
     authenticator = BitbucketOAuthenticator()
-    authenticator.team_whitelist = ['blue']
+    authenticator.bitbucket_team_whitelist = ['blue']
 
     teams = {
         'red': ['grif', 'simmons', 'donut', 'sarge', 'lopez'],
@@ -80,7 +80,3 @@ def list_teams(request):
     handler = client.handler_for_user(user_model('donut'))
     name = yield authenticator.authenticate(handler)
     assert name == 'donut'
-
-
-
-

oauthenticator/tests/test_github.py

@@ -1,4 +1,12 @@
+import re
+import functools
+import json
+from io import BytesIO
+
 from pytest import fixture, mark
+from urllib.parse import urlparse, parse_qs
+from tornado.httpclient import HTTPRequest, HTTPResponse
+from tornado.httputil import HTTPHeaders
 
 from ..github import GitHubOAuthenticator
 
@@ -33,3 +41,84 @@ def test_github(github_client):
 @mark.gen_test
 def test_no_code(github_client):
     yield no_code_test(GitHubOAuthenticator())
+
+
+def make_link_header(urlinfo, page):
+    return {'Link': '<{}://{}{}?page={}>;rel="next"'
+                    .format(urlinfo.scheme, urlinfo.netloc, urlinfo.path, page)}
+
+
+@mark.gen_test
+def test_org_whitelist(github_client):
+    client = github_client
+    authenticator = GitHubOAuthenticator()
+
+    ## Mock Github API
+
+    teams = {
+        'red': ['grif', 'simmons', 'donut', 'sarge', 'lopez'],
+        'blue': ['tucker', 'caboose', 'burns', 'sheila', 'texas'],
+    }
+
+    member_regex = re.compile(r'/orgs/(.*)/members')
+
+    def team_members(paginate, request):
+        urlinfo = urlparse(request.url)
+        team = member_regex.match(urlinfo.path).group(1)
+
+        if team not in teams:
+            return HTTPResponse(400, request)
+
+        if not paginate:
+            return [user_model(m) for m in teams[team]]
+        else:
+            page = parse_qs(urlinfo.query).get('page', ['1'])
+            page = int(page[0])
+            return team_members_paginated(
+                team, page, urlinfo, functools.partial(HTTPResponse, request))
+
+    def team_members_paginated(team, page, urlinfo, response):
+        if page < len(teams[team]):
+            headers = make_link_header(urlinfo, page + 1)
+        elif page == len(teams[team]):
+            headers = {}
+        else:
+            return response(400)
+
+        headers.update({'Content-Type': 'application/json'})
+
+        ret = [user_model(teams[team][page - 1])]
+
+        return response(200,
+                        headers=HTTPHeaders(headers),
+                        buffer=BytesIO(json.dumps(ret).encode('utf-8')))
+
+    ## Perform tests
+
+    for paginate in (False, True):
+        client.hosts['api.github.com'].append(
+            (member_regex, functools.partial(team_members, paginate)),
+        )
+
+        authenticator.github_organization_whitelist = ['blue']
+
+        handler = client.handler_for_user(user_model('caboose'))
+        name = yield authenticator.authenticate(handler)
+        assert name == 'caboose'
+
+        handler = client.handler_for_user(user_model('donut'))
+        name = yield authenticator.authenticate(handler)
+        assert name is None
+
+        # reverse it, just to be safe
+        authenticator.github_organization_whitelist = ['red']
+
+        handler = client.handler_for_user(user_model('caboose'))
+        name = yield authenticator.authenticate(handler)
+        assert name is None
+
+        handler = client.handler_for_user(user_model('donut'))
+        name = yield authenticator.authenticate(handler)
+        assert name == 'donut'
+
+        client.hosts['api.github.com'].pop()