/**
Copyright: 2024 Contrast Security, Inc
  Contact: support@contrastsecurity.com
  License: Commercial

  NOTICE: This Software and the patented inventions embodied within may only be
  used as part of Contrast Security’s commercial offerings. Even though it is
  made available through public repositories, use of this Software is subject to
  the applicable End User Licensing Agreement found at
  https://www.contrastsecurity.com/enduser-terms-0317a or as otherwise agreed
  between Contrast Security and the End User. The Software may not be reverse
  engineered, modified, repackaged, sold, redistributed or otherwise used in a
  way not consistent with the End User License Agreement.
*/
/**
 * @module lib/assess/sinks/mongodb
 *
 * Sinks for nosql injection in mongodb. Collection is where user code calls
 * into whatever it is that we care about. However, there are a lot of methods in here
 * and it's a hassle to keep track of how we need to assess vulnerability in each.
 *
 * Thankfully, Collections's methods dispatch to CoreServer. We use the CoreServer hook to actually
 * assess vulnerability, and hook Collection for the sole purpose of passing the correct (user-relevant)
 * context information.
 */
'use strict';

const _ = require('lodash');
const semver = require('semver');
const logger = require('../../core/logger')('contrast:hooks:mongodb');
const patcher = require('../../hooks/patcher');
const { PATCH_TYPES } = require('../../constants');
const requireHook = require('../../hooks/require');
const { Signature, CallContext } = require('../models');
const policy = require('../policy');
const Scopes = require('../../core/async-storage/scopes');

const ruleId = 'nosql-injection';
const disallowedTags = [
  'limited-chars',
  'alphanum-space-hyphen',
  'string-type-checked',
  'custom-validated-nosql-injection'
];
const requiredTags = ['untrusted'];

module.exports = ({ common }) => {
  const mongoSink = {};

  /**
   * Methods to hook in Collection.prototype
   * Collated and pruned from: Object.keys(mongodb.Collection.prototype)
   * Some methods are deprecated in newer DATABASE versions
   */
  mongoSink.collectionMethods = [
    // XXX hooking these results in compilation issues; unsure why
    // 'collectionName',
    // 'namespace',
    // 'readConcern',
    // 'writeConcern',
    // 'hint',

    // ignoring methods which don't make sense to watch
    // 'drop',
    // 'isCapped',
    // 'listIndexes',
    // 'aggregate',
    // 'geoNear',
    // 'geoHaystackSearch',
    // 'options',
    // 'count',
    // 'stats',
    // 'parallelCollectionScan',
    // 'group',
    // 'insert',
    // 'insertMany',
    // 'insertOne',

    // related to indexes; unsure what, if anything, needs to be done with these
    // 'createIndex',
    // 'createIndexes',
    // 'dropIndex',
    // 'dropIndexes',
    // 'dropAllIndexes',
    // 'reIndex',
    // 'ensureIndex',
    // 'indexExists',
    // 'indexInformation',
    // 'indexes',

    // may need special casing
    // 'distinct',
    // 'mapReduce',
    // 'initializeUnorderedBulkOp',
    // 'initializeOrderedBulkOp'

    // 'bulkWrite',

    'find',
    'findOne',
    'findAndModify',
    'findOneAndDelete',
    'findOneAndReplace',
    'findOneAndUpdate',
    'remove',
    'removeOne',
    'replaceOne',
    'removeMany',
    'save',
    'update',
    'updateOne',
    'updateMany',
    'deleteOne',
    'deleteMany',
  ];

  // Collection of methods with the syntax: method(filter, doc, options);
  // We should track and report unsafe inputs in both: `filter` and `doc`
  const twoArgsCollectionMethods = [
    'findOneAndReplace',
    'findOneAndUpdate',
    'replaceOne',
    'updateOne',
    'updateMany'
  ];

  /**
   * Pre handler for collectionMethods
   */
  mongoSink.collectionMethodHandler = (data) => {
    const doc = data.args[0];

    if (doc && !doc.__contrastContext) {
      try {
        const ctxt = CallContext.create({
          ...data,
          result: null
        });

        // events are prefixed with 'protect '
        ctxt.signature = data.name.split(' ')[1];
        Object.defineProperty(doc, '__contrastContext', {
          value: ctxt,
          enumerable: false
        });

        if (
          Array.isArray(doc) &&
          typeof doc[0] === 'object' &&
          // don't redefine if already present
          !doc[0].__contrastContext
        ) {
          Object.defineProperty(doc[0], '__contrastContext', {
            value: ctxt,
            enumerable: false
          });
        }
      } catch (err) {
        logger.error(`unable to track query in ${data.name} %o`, err);
      }
    }
  };

  mongoSink.collectionMethodPost = (data) => {
    const method = data.hooked && data.hooked.name;

    try {
      let doc = data.args[0];

      if (!doc) {
        return;
      }

      const contrastContext = doc.__contrastContext;

      if (!contrastContext) {
        return;
      }

      const ctxt = getCallContext(contrastContext, data);

      if (twoArgsCollectionMethods.indexOf(method) > -1) {
        // Assess both args at once
        doc = { filter: data.args[0], update: data.args[1] };
      }

      mongoSink.assess(doc, ctxt);

    } catch (err) {
      logger.error(`Unable to perform post hook in ${data.name} %o`, err);
    }
  };

  /**
   * Methods to hook in CoreServer.prototype which take queries
   */
  mongoSink.coreServerQueryMethods = ['command', 'cursor'];

  /**
   * Post handler for coreServerQueryMethods
   */
  mongoSink.coreServerQueryHandler = (version, data) => {
    let query = _.get(data, 'args[1].query');

    // signature changed in 3.3.0 to attach query to
    // a collection class(cmd)
    // also applies to v4 Db.command
    if (semver.gte(version, '3.3.0')) {
      query = _.get(data.args[0], 'cmd.query');
    }

    if (!query) {
      return;
    }
    const contrastContext = _.get(query, '__contrastContext');
    const ctxt = getCallContext(contrastContext, data);
    mongoSink.assess(query, ctxt);
  };

  /**
   * Methods to hook in CoreServer.prototype which take arrays of queries
   * in v4 these methods are available on the `Collection` class.
   */
  mongoSink.coreServerTransactionMethods = ['insert', 'remove', 'update'];

  /**
   * Post handler for coreServerTransactionMethods
   */
  mongoSink.coreServerTransactionHandler = (data) => {
    // FIXME: this is ignoring all but the first query
    // we will eventually want to pick a portion of the input
    // (if array length > X) to sample for finding context/checking
    // for vulnerability
    const cmd = _.get(data, 'args[1][0]');
    if (!cmd) {
      return;
    }

    const contrastContext = _.get(
      cmd,
      '__contrastContext',
      _.get(cmd, 'q.__contrastContext')
    );
    const ctxt = getCallContext(contrastContext, data);
    mongoSink.assess(cmd, ctxt);
  };

  /**
   * Report the vulnerabilikty given the right evidence.
   * FIXME: it's a bit of a hack to use the name to set the signature;
   * if the name scheme changes, we'll report differently
   *
   * @param {string} queryVal The string value in reply which is untrusted
   * @param {CallContext} ctxt The CallContext used in building the event to report
   */
  mongoSink.report = function report(queryVal, ctxt) {
    const sigString = ctxt.signature;
    const sigArray = sigString.split('.');
    const moduleName = sigArray.shift();
    const methodName = sigArray.join('.');
    const signature = new Signature({ moduleName, methodName });
    common.report({ ruleId, signature, input: queryVal, ctxt });
  };

  /**
   * Looks at a query, checks if it is vulnerable, and reports
   *
   * @param {object} query The query object in the Mongo function.
   * @param {CallContext} context Context of the original call in user code.
   * 		Not currently used for checking vulnerability, and just passed through
   * 		to reporting.
   */
  mongoSink.assess = (query, context) => {
    const searchDepth = 3;
    let vulnerableString;

    Scopes.runInAllowAllScope(() => {
      vulnerableString = common.isVulnerable({
        searchDepth,
        disallowedTags,
        requiredTags,
        input: query
      });
    }, 'mongodbSink.assess');

    if (vulnerableString) {
      mongoSink.report(vulnerableString, context);
    }
  };

  mongoSink.handleVersion4 = (mongodb, version) => {
    patcher.patch(mongodb.Db.prototype, 'command', {
      name: 'assess mongodb.Db.prototype',
      patchType: PATCH_TYPES.ASSESS_SINK,
      post: mongoSink.coreServerQueryHandler.bind(null, version)
    });
  };

  mongoSink.handleVersion3 = (mongodb, version) => {
    patcher.patch(
      mongodb.CoreServer.prototype,
      mongoSink.coreServerQueryMethods,
      {
        name: 'assess mongodb.CoreServer.prototype',
        patchType: PATCH_TYPES.ASSESS_SINK,
        post: mongoSink.coreServerQueryHandler.bind(null, version)
      }
    );

    patcher.patch(
      mongodb.CoreServer.prototype,
      mongoSink.coreServerTransactionMethods,
      {
        name: 'assess mongodb.CoreServer.prototype',
        patchType: PATCH_TYPES.ASSESS_SINK,
        post: mongoSink.coreServerTransactionHandler
      }
    );
  };

  /**
   * Installs instrumentation
   */
  mongoSink.handle = () => {
    const enabled =
      _.get(policy, 'rules.nosql-injection-mongodb.enabled') &&
      _.get(policy, 'rules.nosql-injection.enabled');

    if (!enabled) {
      return;
    }

    // NPM package architectures are different between v3.x.x and v4.x.x
    requireHook.resolve({ name: 'mongodb' }, function(mongodb, { version }) {
      patcher.patch(mongodb.Collection.prototype, mongoSink.collectionMethods, {
        name: 'assess mongodb.Collection.prototype',
        patchType: PATCH_TYPES.ASSESS_SINK,
        pre: mongoSink.collectionMethodHandler,
        post: mongoSink.collectionMethodPost
      });

      if (semver.gte(version, '4.0.0')) {
        return mongoSink.handleVersion4(mongodb, version);
      }

      return mongoSink.handleVersion3(mongodb, version);
    });
  };

  return mongoSink;
};

/**
 * It tries to get the context from the pre handler
 * See: mongoSink.collectionMethodHandler
 *
 * If it does not exist, it tries to create a call context
 */
function getCallContext(ctxt, data) {
  if (!ctxt) {
    ctxt = CallContext.create(data);
    ctxt.signature = data.name;
  }
  ctxt.result = data.result;
  return ctxt;
}

module.exports.disallowedTags = disallowedTags;
