Auth0 rule to retrieve Azure AD groups on login

azure-ad-groups-with-client-credentials.js

// This rule will get the groups for users coming from Azure AD
// Auth0 already has the option to do that, but it (currently) won't work
// if the user is coming from a different directory than the directory
// where the app is registered (this can happen with multi-tenant apps).
// This is a variation that gets an access token for Azure AD using the
// client-credential grants instead of using the access token given to the user.
// It's useful if a new access token from Azure AD is not obtained every time the rule runs,
// or if WS-Federation is used instead of OIDC.
// 
// After the rule runs, you will have the 'groups' property in the user
// that you can use to add custom claims to the id_token.
// 
// Currently, the rule just reads the top-level groups:
//       waadClient.getGroupsForUserByObjectIdOrUpn(user.oid, function(err, groups) {
// but you can change it to read the extended user profile:
//       waadClient.getUserByProperty('objectId', user.oid, { includeGroups: true, includeNestedGroups: false }, function(err, extUser) {
//
// To set up:
//   Configure the Azure AD domain, application ID and Key (i.e. client id and secrets, from the WAAD connection)
//   You can use the [configuration object](https://auth0.com/docs/rules/current#using-the-configuration-object)
//   to store those values.
//   The application registration in Azure AD needs to have the appropriate permissions to read users and groups (this is an app permission, not a delegated permission).
//   This rules assumes just one Azure AD connection, otherwise you'll have to use different ids/secret for each Azure AD domain.


function (user, context, callback) {
  'use strict';
  if (context.connectionStrategy !== 'waad') {
    // don't execute for non Azure AD connections
    return callback(null, user, context);
  }
  if (user.groups) {
    // dont't execute if user groups were already retrieved
    return callback(null, user, context);
  }

  var WAAD_DOMAIN = 'the azure AD domain';
  var WAAD_APP_ID = 'the azure AD app ID';
  var WAAD_APP_KEY = 'the azure AD app key';
  
  var request = require('request');
  var async = require('async');
  function Waad10(options) {
    options = options || {};
    this.options = options;
    this.baseUrl = 'https://graph.windows.net/';
    this.apiVersion = '1.0';

    if (!this.options.tenant) {
      throw new Error('Must supply "tenant" id (16a88858-..2d0263900406) or domain (mycompany.onmicrosoft.com)');
    }

    if (!this.options.accessToken) {
      throw new Error('Must supply "accessToken"');
    }
  }
  Waad10.prototype.getUserByProperty = function (propertyName, propertyValue, options, callback) {
    if (typeof options === 'function') {
      callback = options;
      options = {};
    }

    var qs = {
      "$filter": propertyName + " eq '" + propertyValue + "'",
      "$top": 1
    };

    return this.__queryUsers(qs, options, function (err, users) {
      if (err) return callback(err);
      return callback(null, users[0]);
    });
  };
  Waad10.prototype.__queryUsers = function (qs, options, callback) {
    if (typeof options === 'function') {
      callback = options;
      options = options || {};
    }

    qs['api-version'] = this.apiVersion;

    this.__request({
      url: this.baseUrl + this.options.tenant + '/users',
      qs: qs,
    }, function (err, users) {
      if (err || !users) return callback(err);

      if (!options || !options.includeGroups) return callback(null, users);


      async.forEach(users, function (user, cb) {
        var groupsCallback = function (err, groups) {
          if (err) return callback(err);
          user.groups = groups;
          cb();
        };

        if (options.includeNestedGroups) {
          this.__queryNestedUserGroups(user.objectId, groupsCallback);
        } else {
          this.__queryUserGroup(user.objectId, groupsCallback);
        }
      }.bind(this), function (err) {
        if (err) return callback(err);
        return callback(null, users);
      });
    }.bind(this));
  };
  Waad10.prototype.__queryUserGroup = function (objectIdOrUpn, cb) {
    var qs = {
      "api-version": this.apiVersion
      // Uncomment if you want to change the default max of 100 groups returned
      //,"$top": 250
    };
    this.__request({
      url: this.baseUrl + this.options.tenant + '/users/' + objectIdOrUpn + '/memberOf',
      qs: qs
    }, cb);
  };
  Waad10.prototype.__queryNestedUserGroups = function (objectIdOrUpn, cb) {
    var qs = {
      "api-version": this.apiVersion
      // Uncomment if you want to change the default max of 100 groups returned
      //,"$top": 250
    };

    this.__request({
      url: this.baseUrl + this.options.tenant + '/users/' + objectIdOrUpn + '/getMemberGroups',
      qs: qs,
      method: 'POST',
      json: { "securityEnabledOnly": false }
    }, function (err, groups) {
      if (err || !groups) return cb(err);

      return this.__getObjectsByObjectIds(groups, cb);
    }.bind(this));
  };

  Waad10.prototype.__getObjectsByObjectIds = function (objectIds, cb) {
    var qs = {};
    qs['api-version'] = '1.6';

    this.__request({
      url: this.baseUrl + this.options.tenant + '/getObjectsByObjectIds',
      qs: qs,
      method: 'POST',
      json: {
        "objectIds": objectIds,
        "types": ["group"]
      }
    }, cb);
  };
  Waad10.prototype.__request = function (options, callback) {
    var headers = {
      'Authorization': 'Bearer ' + this.options.accessToken,
    };

    request({
      url: options.url,
      qs: options.qs || {},
      method: options.method || 'GET',
      json : options.json,
      headers: headers
    }, function(err, resp, body) {
      if (err) return callback(err, null);

      if (resp.statusCode !== 200) {
        return callback(new Error(body), null);
      }

      var array = body;

      if (typeof body === 'string'){
        try {
          array = JSON.parse(body);
        } catch (e){
          return callback(new Error('Failed to parse response from graph API'));
        }
      }

      if (!array && array.value && array.value.length === 0)
        return callback();

      return callback(null, array.value);
    }.bind(this));
  };
  Waad10.prototype.getGroupsForUserByObjectIdOrUpn = function(objectIdOrUpn, callback) {
    return this.__queryUserGroup(objectIdOrUpn, callback);
  };  

  function getWaadAccessTokenWithClientCredentials2(tenantDomain, clientId, clientSecret, callback) {
    var cacheKey = "WaadTokenToReadGroups";
    var cachedToken = global[cacheKey];
    if (cachedToken) {
      if (cachedToken.expirationDate > new Date()) {
        console.log("Reusing cached token that expires in "+ cachedToken.expirationDate);
        return callback(null, cachedToken.token);
      }
      console.log('Cached token found but expired.');
    } else {
      console.log('Cached token not found.');
    }

    var data = {
      grant_type: 'client_credentials',
      client_id: clientId,
      client_secret: clientSecret,
      resource: '00000002-0000-0000-c000-000000000000/graph.windows.net@' + tenantDomain
    };
  
    request.post('https://login.windows.net/' + tenantDomain + '/oauth2/token', { form: data }, function(e, resp, body) {
      if (e) return callback(e, null);
      var response;
      if (resp.statusCode !== 200) {
        try {
          response = JSON.parse(body);
          if (response.error) {
            return callback(response.error_description || response.error);
          }
        }
        catch (exp) {}
  
        return callback(body);
      }
      response = JSON.parse(body);

      var expirationDate = new Date();
      expirationDate.setSeconds(expirationDate.getSeconds() + response.expires_in - 60);      
      // store token in cache
      console.log("Storing token, expires in "+ expirationDate);
      global[cacheKey] = {
        expirationDate: expirationDate,
        token: response.access_token
      };

      callback(null, response.access_token);
    });
  }
  
  getWaadAccessTokenWithClientCredentials2(WAAD_DOMAIN, WAAD_APP_ID, WAAD_APP_KEY, function (err, waadToken) {
    if (err) {
      return callback(err);
    }
    var waadClient = new Waad10({ accessToken: waadToken, tenant: user.tenantid });
    waadClient.getGroupsForUserByObjectIdOrUpn(user.oid, function(err, groups) {
      if (err) {
        return callback(err);
      }
      user.groups = [];
      if (groups) {
        groups.forEach(function (group) {
          user.groups.push(group.displayName);
        });
      }
      callback(null, user, context);
    });
  });
}