Switched herokai_only flag over to a general purpose `auth_callback…

…` instead.

README.md

@@ -47,6 +47,10 @@ you're using. For example, in Django you'll find this in
 
 This is a dict-like object with a couple of useful keys:
 
+* `"user"` - the Heroku 
+  [account object](https://devcenter.heroku.com/articles/platform-api-reference#account) 
+  from the Platform API.
+
 * `"username"` - the Heroku account email address (same as `env["REMOTE_USER"]`).
 
 * `"access_token"` - the OAuth user access token. You can use this to
@@ -91,9 +95,17 @@ Those options are:
   There's not a great reason to set this to `False`, but you can if you
   really feel like it I guess.
 
-* `herokai_only` - `False` by default. If `True`, only allow logins from
-  `@heroku.com` email addresses. You probably don't want this unless you
-  work at Heroku.
+* `auth_callback` - an optional callback function that will be called to decide
+  if the authorized user should be allowed in. This will be passed a single
+  arguent, the [session object](#session). For example, to limit access to
+  users from a single domain:
+
+```python
+def callback(session):
+    return session['user']['email'].endswith('@example.com')
+
+app = bouncer(app, auth_callback=callback)
+```
 
 * `scope` - the OAuth scope(s), as [defined in Heroku's documentation](https://devcenter.heroku.com/articles/oauth#scopes),
   that your app requires. If you're requesting more than one scope, the scopes

heroku_bouncer.py

@@ -5,15 +5,21 @@
 def bouncer(application, set_remote_user=True, herokai_only=False,
             path='/auth/heroku/callback/', cookie='herokuoauthsess',
             forbidden_path='/auth/heroku/forbidden/', forbidden_passthrough=False,
-            client_id=None, client_secret=None, secret_key=None, scope='identity'):
+            client_id=None, client_secret=None, secret_key=None, scope='identity',
+            auth_callback=None):
     try:
         client_id = client_id or os.environ['HEROKU_OAUTH_ID']
         client_secret = client_secret or os.environ['HEROKU_OAUTH_SECRET']
         secret_key = secret_key or os.environ['SECRET_KEY']
     except KeyError as e:
         raise EnvironmentError("Missing configuration in environ: %s" % e)
 
-    service = HerokuService(herokai_only=herokai_only)
+    if herokai_only:
+        import warnings
+        warnings.warn("herokai_only is deprecated; use auth_callback instead", DeprecationWarning)
+        auth_callback = lambda s: s["user"]["email"].endswith("@heroku.com")
+
+    service = HerokuService(auth_callback=auth_callback)
     client = service.make_client(client_id=client_id, client_secret=client_secret, scope=scope)
     return client.wsgi_middleware(
         application = application,
@@ -26,20 +32,21 @@ def bouncer(application, set_remote_user=True, herokai_only=False,
     )
 
 class HerokuService(wsgioauth2.Service):
-    def __init__(self, herokai_only=False):
+    def __init__(self, auth_callback=None):
         super(HerokuService, self).__init__(
             authorize_endpoint="https://id.heroku.com/oauth/authorize",
             access_token_endpoint="https://id.heroku.com/oauth/token")
-        self.herokai_only = herokai_only
+        self.auth_callback = auth_callback
 
     def load_username(self, access_token):
         headers = {'Authorization': 'Bearer %s' % access_token,
                    'Accept': 'application/vnd.heroku+json; version=3'}
         resp = requests.get('https://api.heroku.com/account', headers=headers)
-        access_token['username'] = resp.json()['email']
+        access_token['user'] = resp.json()
+        access_token['username'] = access_token['user']['email']
 
     def is_user_allowed(self, access_token):
-        if self.herokai_only:
-            return access_token['username'].endswith('@heroku.com')
+        if self.auth_callback:
+            return self.auth_callback(access_token)
         else:
             return True