/**
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.
*/
'use strict';
const patcher = require('../../hooks/patcher');
const AsyncStorage = require('../../core/async-storage');
const { PATCH_TYPES } = require('../../constants');
const moduleHook = require('../../hooks/require');
const _ = require('lodash');
const { Signature, CallContext } = require('../models');
const ruleId = 'nosql-injection-dynamodb';
const disallowedTags = [
  'limited-chars',
  'alphanum-space-hyphen',
  'string-type-checked',
];
const requiredTags = ['untrusted'];

/*
*  Schema of attributes the Node Agent is looking for user-controlled data
*  The only exception to the rule so far is the ScanFilter in ScanCommand.
*  It has nested schema too. So far we are only reporting a NoSQL Injection
*  if ComparisonOperator within the ScanFilter is user-controlled.
*  ScanFilter is handled individually in extractValues function
* */
const trackSchemaCommands = {
  'scan': {
    attributes: [
      'FilterExpression',
      'ExclusiveStartKey',
      'ScanFilter'
    ]
  },
  'executeStatement': { attributes: ['Statement'] },
  'ScanCommand': {
    attributes: [
      'ComparisonOperator',
      'FilterExpression',
      'ProjectionExpression',
      'ScanFilter',
    ]
  },
  'ExecuteStatementCommand': {
    attributes: ['Statement']
  }
};

const requests = new WeakSet();

/**
 * Extracts all values from either a dynamo document client or client
 *
 * @param {Object} payload dynamo scan payload
 * @param {string} mode client, docClient or command (for aws sdk v3)
 * @param {string} version DynamoDB SDK version
 * @return {Array} all string values from payload
 */
function extractValues(command, payload = {}) {
  return _.flatten(
    trackSchemaCommands[command].attributes.map((key) => {
      if (payload[key] == undefined) return;
      if (typeof payload[key] === 'string') return payload[key];

      // ScanFilter is an exception. It is almost safe for any nested attribute
      if (key === 'ScanFilter') {
        // collect the values from ComparisonOperator attributes ONLY
        return getComparisonValues(payload[key]);
      }

      const values = _.values(payload[key]);
      const extractionMethod = findValues.bind(this);
      return _.flattenDeep(extractionMethod(values));
    }),
  );
}

/**
 * We only track strings. some data types
 * can contain strings as keys or values but
 * null, Buffers and Numbers are objects
 * but do not have trackable values so skip these
 *
 * @param {any} value check if trackable
 * @return {boolean}
 */
function notTrackable(value) {
  return value === null || Buffer.isBuffer(value) || _.isNumber(value);
}

/**
 * Extracts all strings from a dynamo document client payload
 *
 * Note: There are other data types dynamo supports
 * to cast to AttributeValue types, but we currently
 * do not track null, Boolean, Number, Buffer(a few types here)
 *
 * @param {Object} values for all keys in relevant properties
 * @return {Array} all values
 */
function findValues(values) {
  return _.map(values, (value) => {
    switch (typeof value) {
      case 'string':
        return value;
      case 'object': {
        if (notTrackable(value)) {
          return;
        }
        const values = _.values(value);
        return findValues(values);
      }
    }
  });
}

/**
 * Check if key exists in ScanFilter bject and return the value if so
 * Each element of ScanFilter is an object with predictable structure
 * "Genre": {
 *   "AttributeValueList":[ {"S":"Rock"} ],
 *   "ComparisonOperator": "EQ"
 * }
 *
 * @param {Object} values for all keys in a given payload
 * @param {String} key name we are looking the value of
 * @return {Any} if value is found for a given key
 */
const getComparisonValues = (obj) => Object.keys(obj).map(field => {
  if (typeof obj[field] === 'object' && Object.prototype.hasOwnProperty.call(
    obj[field],
    'ComparisonOperator'
  )) {
    return obj[field].ComparisonOperator;
  }
});

module.exports = ({ common }) => {
  const { isVulnerable, report } = common;

  const dynamoSink = {};
  /**
   * Registers the hooks for client and document client scan
   */
  dynamoSink.handle = function () {
    moduleHook.resolve({ name: 'aws-sdk' }, (aws) => {
      const client = aws.DynamoDB.prototype;

      patcher.patch(client, 'makeRequest', {
        name: 'aws-sdk.DynamoDB.prototype',
        patchType: PATCH_TYPES.ASSESS_SINK,
        alwaysRun: true,
        post(data) {
          // FIXME NODE-1337 remove this check and remove alwaysRun
          // option once that has been done.
          // We need to alwaysRun this patch currently b/c we also
          // patch this function to report dynamo as a component for
          // flowmap, which has nothing to do w/ request scope.
          if (!AsyncStorage.getContext()) {
            return;
          }

          if (Object.keys(trackSchemaCommands).includes(data.args[0]) && data.args[1]) {
            if (requests.has(data.args[1])) {
              return;
            }

            // since this is instrumenting a wrapper we only want to say the args
            // were the 2nd arg which is the dynamo db payload
            const ctxtData = {
              args: [data.args[1]],
              obj: data.obj,
              result: data.result,
            };

            const values = extractValues(data.args[0], data.args[1]);
            dynamoSink.check({
              values,
              methodName: `DynamoDB.prototype.${data.args[0]}`,
              data: ctxtData,
            });
          }
        },
      });

      // DocumentClient added in aws-sdk 2.2.
      if (aws.DynamoDB.DocumentClient) {
        const docClient = aws.DynamoDB.DocumentClient.prototype;
        /**
         * We add request params in a WeakSet so we can check in makeRequest above and not
         * execute instrumentation. This library has an interesting architecture and we want
         * the right instrumentation to run for requests
         */
        patcher.patch(docClient, 'makeServiceRequest', {
          name: 'aws-sdk.DynamoDB.DocumentClient.prototype',
          patchType: PATCH_TYPES.ASSESS_SINK,
          pre(data) {
            if (data.args[0] === 'scan' && data.args[1]) {
              requests.add(data.args[1]);
            }
          },
        });

        patcher.patch(docClient, 'scan', {
          name: 'aws-sdk.DynamoDB.DocumentClient.prototype',
          patchType: PATCH_TYPES.ASSESS_SINK,
          post(data) {
            const values = extractValues('scan', data.args[0]);
            dynamoSink.check({
              values,
              methodName: 'DynamoDB.DocumentClient.prototype.scan',
              data,
            });
          },
        });
      }
    });

    moduleHook.resolve({ name: '@aws-sdk/client-dynamodb' }, (aws) => {
      const client = aws.DynamoDBClient.prototype;

      patcher.patch(client, 'send', {
        name: '@aws-sdk/client-dynamodb.ScanCommand.prototype',
        patchType: PATCH_TYPES.ASSESS_SINK,
        alwaysRun: true,
        post(data) {
          if (!AsyncStorage.getContext()) return;

          if (data.args[0] && data.args[0].constructor && data.args[0].input) {
            const sendCommand = data.args[0].constructor.name;

            if (Object.keys(trackSchemaCommands).includes(sendCommand)) {
              const values = extractValues(
                sendCommand,
                data.args[0].input
              ).filter(Boolean);

              dynamoSink.check({
                data,
                values,
                methodName: `DynamoDBClient.${sendCommand}`,
              });
            }
          }
        },
      });
    });

    return dynamoSink;
  };

  dynamoSink.check = function ({ data, values, methodName }) {
    const vulnerableString = _.compact(
      values.map((input) =>
        isVulnerable({ disallowedTags, requiredTags, input }),
      ),
    );

    if (vulnerableString.length) {
      const ctxt = CallContext.create(data);
      const signature = new Signature({ moduleName: 'aws-sdk', methodName });
      report({ ruleId, signature, input: vulnerableString[0], ctxt });
    }
  };

  return dynamoSink;
};
