/**
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.
*/
/**
 * String.prototype.replace Propagation Provider module.
 * @module lib/assess/propagators/String.prototype.replace
 */
'use strict';

const _ = require('lodash');

const TagRange = require('../models/tag-range');
const tagRangeUtil = require('../models/tag-range/util.js');
const tracker = require('../../tracker.js');
const { PropagationEvent, Signature, CallContext } = require('../models');
const Scopes = require('../../core/async-storage/scopes');
const { unwrap } = require('../../hooks/patcher');
const logger = require('../../core/logger')(
  'contrast:rules:propagation:String.prototype.replace'
);

const PROPAGATOR_METHOD = 'String.prototype.replace';

module.exports.handle = {
  pre(callContext) {
    // In the event of our propagation causing an unexpected result,
    // we want to make sure that we have a safety net given how much complicated
    // logic we go through to propagate.
    // For an example of how this breaks, view NODE-1190
    // we only want to run this in the event that the user is not sending in a custom function
    // because there's no generic way to ensure that running String.prototype.replace twice
    // with a custom function will not incur side effects.
    if (!usesCustomReplacer(callContext.args)) {
      const origReplace = unwrap(String.prototype.replace);
      callContext._originalResult = origReplace.apply(
        callContext.obj,
        callContext.args
      );
    }

    const noop = true;
    const trackData = getTrackData(callContext);

    coerceEmptyReplacerArgs(callContext);
    callContext._state = {
      trackData,
      orig: {
        obj: callContext.obj,
        args: [...callContext.args]
      },
      noop
    };

    // bail if propagation is n/a
    if (
      !trackData.objData.tracked &&
      !trackData.replacerData.tracked &&
      typeof callContext.args[1] != 'function'
    ) {
      return;
    }

    // otherwise instrument via wrapped replacer function
    callContext.args[1] = createCustomReplacerSpy({ callContext, trackData });
  },
  post(callContext) {
    if (callContext._state.noop) {
      return;
    }

    // in the event of failed propagation, restore the result to be the result of the original
    // replace and log in the event of false negatives.
    if (
      callContext._originalResult &&
      callContext._originalResult != callContext.result
    ) {
      logger.warn(
        `String replace result does not match expected value. Real: ${callContext.result}, Expected: ${callContext._originalResult}`
      );
      callContext.result = callContext._originalResult;
    } else {
      trackFinalResult(callContext);
    }
  }
};

/**
 * Initializes and returns the state model used for calculations throughout.
 * @param {Object} callContext
 * @returns {Object}
 */

function getTrackData(callContext) {
  // The default data template; logic below will modify template accordingly.
  const trackData = {
    objData: { tracked: false },
    replacerData: { tracked: false },
    dynamicTagRanges: [],
    dynamicEvents: new Set(),
    substitutionDelta: 0,
    totalShift: 0,
    resultTagRanges: []
  };

  const objData = tracker.getData(callContext.obj);
  if (objData) {
    trackData.objData = objData;
    trackData.resultTagRanges = objData.tagRanges.map((r) => r.clone());
  }

  const replacerData = tracker.getData(callContext.args[1]);
  if (replacerData) {
    trackData.replacerData = replacerData;
  }

  return trackData;
}

/**
 * Check if the String.replace call uses a custom replacer function
 * @param   {*[]}   args callContext.args array for call
 * @returns boolean      is arg[1]or arg[2] a function
 */
function usesCustomReplacer(args) {
  return typeof args[1] === 'function' || typeof args[2] === 'function';
}

/**
 * Coerces `undefined` or `null` into the string representation. Otherwise will
 * return the original value.
 * @param {*} replacer
 * @returns {string}
 */
function coerceReplacer(replacer) {
  if (replacer === undefined) {
    replacer = 'undefined';
  } else if (replacer === null) {
    replacer = 'null';
  }
  return replacer;
}

/**
 * Mutates the call context to swap replacer argument w/ its coerced value.
 * @param {Object} callContext
 */
function coerceEmptyReplacerArgs(callContext) {
  callContext.args[1] = coerceReplacer(callContext.args[1]);
  return callContext;
}

/**
 * Given a replacer function's arguments will return the match and offset
 * values.
 * @param {Array} args
 * @returns {Array}
 */
function getMatchAndOffsetValues(args) {
  const [match] = args;
  const offset = findOffset(args);
  return [match, offset];
}

/**
 * Looks through a replacer's argument array and returns the 1st value which is
 * a number.
 * @param {Array} args
 * @returns {Number}
 */
function findOffset(args) {
  return _.isNumber(args[1])
    ? args[1] // no general capture groups
    : _.isNumber(args[args.length - 2])
      ? args[args.length - 2] // no named-capture groups
      : _.isNumber(args[args.length - 3])
        ? args[args.length - 3] // w/ named-capture groups
        : (() => {
        // Invalid call at this point though
        for (const arg of args) {
          if (_.isNumber(arg)) return arg;
        }
      })();
}

/**
 * Returns the custom replacer function that contrast will set as the 2nd
 * argument to all `string.replace` calls involving tracked data.
 * @param {Object} callContext
 * @returns {Function}
 */
function createCustomReplacerSpy({ callContext }) {
  return function(...args) {
    _.set(callContext, '_state.noop', false);
    const replacement = getReplacementAndCaptureDynamicTags({
      callContext,
      args
    });
    collectRangesForActiveReplacement({ callContext, args, replacement });
    return replacement;
  };
}

/**
 * This will return the replacer string value to be used for the current
 * replacement. Cases for replacer value:
 * - Function
 *   If the replacer is a function, then we wrap it's arguments such that the
 *   values are tracked. If return values are tracked, we will store their tag
 *   ranges in `trackData.dynamicTagRanges`.
 * - String
 *   If the replacer is a string, we manually perform the interpolation for any
 *   of the special character replacements. If any occur, we will track these
 *   new ranges in `trackData.dynamicTagRanges`.
 * @param {object} callContext State
 * @param {*[]} args
 * @returns {string}
 */
function getReplacementAndCaptureDynamicTags({ callContext, args }) {
  const {
    _state: {
      trackData,
      orig: {
        args: [, replacer]
      }
    }
  } = callContext;

  let _replacer;

  if (_.isFunction(replacer)) {
    // Allow propagation within replacer function so we can track return value.
    _replacer = invokeReplacer({ callContext, args, replacer });
    const data = tracker.getData(_replacer);
    if (data) {
      trackData.dynamicTagRanges = data.tagRanges;
      trackData.dynamicEvents.add(data.event);
    } else {
      trackData.dynamicTagRanges = [];
    }
  } else {
    _replacer = getReplacementAndCaptureDynamicStringReplacerTags({
      callContext,
      args
    });
  }

  return _replacer;
}

/**
 * Call replacer within an ALLOW_ALL scope so that we can perform dataflow.
 * @param {object} callContext
 * @param {*[]} args
 * @param {function} replacer
 * @returns {string}
 */
function invokeReplacer({ callContext, args, replacer }) {
  const wrappedArgs = wrapOriginalReplacerArguments({ callContext, args });
  return Scopes.runInAllowAllScope(() => replacer.call(global, ...wrappedArgs));
}

/**
 * When replacer argument is a function, we want to track the arguments passed
 * to it so propagation will occur and thus the result will be tracked.
 * @param {object} callContext
 * @param {*[]} args
 * @returns {*[]}
 */
function wrapOriginalReplacerArguments({ callContext, args }) {
  const {
    _state: {
      orig: {
        args: { [1]: replacer }
      },
      trackData: {
        objData: { tagRanges, tracked }
      }
    }
  } = callContext;

  if (!tracked) {
    return args;
  }

  // Only wrap arguments that are used by replacer,
  for (let idx = 0; idx < replacer.length; idx++) {
    // and only track if needed
    if (tracker.getData(args[idx])) {
      continue;
    }
    trackReplacerArgAtPath({
      callContext,
      target: args,
      path: idx,
      tagRanges
    });
  }

  return args;
}

/**
 * Right now this just tracks full arg values, not accounting for tagRanges
 * calculated from the args' actual positioning in string. This is an accuracy
 * issue that we can improve in the future.
 * @param {object} callContext State
 * @param {*} target Instance having value at `path`
 * @param {number} path Path the the value on `target`
 * @param {object[]} tagRanges Coll. of TagRanges to pass from tracked string
 */
function trackReplacerArgAtPath({ callContext, target, path, tagRanges }) {
  const captured = new Set();
  const _result = _.get(target, path);

  if (_.isString(_result) && _result.length) {
    for (const { tag } of tagRanges) {
      if (captured.has(tag)) {
        continue;
      }

      const tracked = tracker.track(_result);
      if (tracked) {
        // TODO: Improve accuracy here by deriving actual ranges from source,
        // rather than applying "blanket" ranges the length of argument.
        const tagRange = new TagRange(0, tracked.str.length - 1, tag);
        tracked.props.tagRanges.push(tagRange);
        tracked.props.event = createPropagationEvent({
          callContext,
          tagRanges: [tagRange],
          result: tracked.str
        });

        captured.add(tag);
        target[path] = tracked.str;
      }
    }
  } else if (_.isObject(_result)) {
    for (const path of Object.keys(_result)) {
      trackReplacerArgAtPath({ callContext, tagRanges, path, target: _result });
    }
  }
}

/**
 * Returns unique tags from a collection of ranges.
 * @param {TagRange[]} tagRanges
 * @returns {string[]}
 */
function getTagRangesTagsList(tagRanges) {
  return Array.from(
    tagRanges.reduce((acc, { tag }) => acc.add(tag), new Set())
  );
}

/**
 * When the replacer is a string, interpolation can happen when the replacer
 * contains special chars: $$, $&, $', $`. We carry out these replacements in
 * the replacer and collect ranges in `trackData.dynamicTagRanges` when they
 * occur. The current implementation will not take into account the positioning
 * of the interpolated string in the original - we simply create a single tag
 * range that will span the length of the dynamic replacer for each tag type,
 * e.g. "untrusted", on the original string.
 * @param {object} callContext State
 * @param {*[]} args Custom replacer function args
 * @returns {string}
 */
function getReplacementAndCaptureDynamicStringReplacerTags({
  callContext,
  args
}) {
  const {
    _state: {
      trackData,
      orig: {
        obj,
        args: [, replacer]
      }
    }
  } = callContext;
  const [match, offset] = getMatchAndOffsetValues(args);
  let didSpecialReplace = false;

  trackData.dynamicTagRanges = [];

  const replacerRegex =
    callContext.args[0] instanceof RegExp
      ? /\$([$&`']|\d{1,2})/g
      : /\$([$&`'])/g;

  let specialReplacer = replacer
    .toString() // for custom .toString() on objects
    .replace(replacerRegex, (_match, _p1, _offset, _string) => {
      switch (_p1) {
        case '$':
          return '$';
        case '`': {
          didSpecialReplace = true;
          return obj.substring(0, offset);
        }
        case "'": {
          didSpecialReplace = true;
          return obj.substring(offset + match.length);
        }

        case '&': {
          didSpecialReplace = true;
          return match;
        }
        default: {
          const m = args[parseInt(_p1)];
          // There could be empty value corresponding to capture
          if (m) {
            didSpecialReplace = true;
          }
          return m;
        }
      }
    });

  // When using named capture groups in the `search` RegExp (Node >= v12), the
  // last argument is an Object w/ group-name keys
  if (_.isObject(args[args.length - 1])) {
    const groups = args[args.length - 1];
    for (const name of Object.keys(groups)) {
      const value = groups[name];
      specialReplacer = specialReplacer.replace(
        new RegExp(`\\$<${name}>`),
        () => {
          didSpecialReplace = true;
          return value;
        }
      );
    }
  }

  if (didSpecialReplace && trackData.objData.tracked) {
    const tags = getTagRangesTagsList(trackData.objData.tagRanges);
    for (const tag of tags) {
      trackData.dynamicTagRanges = [
        new TagRange(0, specialReplacer.length - 1, tag)
      ];
    }
  }

  return specialReplacer;
}

/**
 * When called, we should have already captured any dynamic tag ranges that
 * result from special pattern interpolation or from user's replacer function
 * return value. At this point we need to remove and adjust tag ranges in the
 * final ranges to account for the portion of the string being replaced, and
 * collect any dynamic or replacer tags and sandwich them into the result also.
 * @param {object} callContext State
 * @param {*[]} args Custom replacer args providing instrumentation context
 * @param {string} replacement
 */
function collectRangesForActiveReplacement({ callContext, args, replacement }) {
  const {
    _state: { trackData }
  } = callContext;
  const [match, offset] = getMatchAndOffsetValues(args);
  const currentShift = coerceReplacer(replacement).length - match.length;

  const { beforeOffset, afterReplacement } = adjustResultTagRangesForReplace({
    tagRanges: trackData.resultTagRanges,
    offset: offset + trackData.totalShift,
    matchLength: match.length,
    currentShift
  });
  const atOffset = getReplacementTagRangesAtOffset(trackData, offset);

  trackData.resultTagRanges = [];
  tagRangeUtil.addAllInPlace(trackData.resultTagRanges, [
    ...beforeOffset,
    ...atOffset,
    ...afterReplacement
  ]);

  trackData.totalShift += currentShift;
}

/**
 * TagRange's for replacement string can come from string replacer values or
 * from dynamic values from replacer functions or custom replacer char'
 * replacements.
 * @param {object} trackData
 * @param {string} offset
 * @returns {TagRange[]}
 */
function getReplacementTagRangesAtOffset(trackData, offset) {
  const { replacerData, dynamicTagRanges, totalShift } = trackData;
  const adjustment = offset + totalShift;
  let atOffset = [];

  if (dynamicTagRanges.length) {
    atOffset = dynamicTagRanges.map(
      (tr) => new TagRange(tr.start + adjustment, tr.stop + adjustment, tr.tag)
    );
  } else if (replacerData.tracked) {
    atOffset = replacerData.tagRanges.map(
      (tr) => new TagRange(tr.start + adjustment, tr.stop + adjustment, tr.tag)
    );
  }
  return atOffset;
}

/**
 * Returns the TagRange[] grouped into "before" and "after" collections around
 * the portion of the string being replaced. This runs with each replacement
 * that occurs in the string.
 * @returns {object} { beforeOffset: [], afterReplacement: [] }
 */
function adjustResultTagRangesForReplace({
  tagRanges,
  offset,
  matchLength,
  currentShift
}) {
  const adjusted = tagRanges.reduce(
    (acc, range) => {
      const { beforeOffset, afterReplacement } = acc;
      const { start, stop, tag } = range;
      if (stop < offset) {
        // entirely before substitution
        beforeOffset.push(range);
      } else if (start >= offset + matchLength) {
        // entirely after substitution
        range.start += currentShift;
        range.stop += currentShift;
        afterReplacement.push(range);
      } else {
        //overlap cases
        if (offset > start) {
          beforeOffset.push(new TagRange(start, offset - 1, tag));
        }
        // add a range from match + adjustedOffset -> stop
        if (stop >= offset + matchLength) {
          afterReplacement.push(
            new TagRange(
              offset + matchLength + currentShift,
              stop + currentShift,
              tag
            )
          );
        }
      }
      return acc;
    },
    {
      beforeOffset: [],
      afterReplacement: []
    }
  );

  return adjusted;
}

/**
 * This gets called in the post hook. Here we track the result, do the final
 * adjustments and merging of the tag ranges tracked throughout the function
 * call. All metadata is added to the tracking data for the result.
 * @param {object} callContext State
 */
function trackFinalResult(callContext) {
  const {
    _state: { trackData }
  } = callContext;

  if (!callContext.result || !trackData.resultTagRanges.length) {
    return;
  }

  if (tracker.getData(callContext.result)) {
    return;
  }

  const tracked = tracker.track(callContext.result);

  if (tracked) {
    tracked.props.event = createPropagationEvent({
      callContext,
      tagRanges: trackData.resultTagRanges,
      result: tracked.str
    });
    tracked.props.tagRanges = trackData.resultTagRanges;
    callContext.result = tracked.str;
  }
}

/**
 * Creates the propagation event used in tracking data.
 * @param {Object} callContext State
 * @param {string} result The tracked value
 * @returns {PropagationEvent}
 */
function createPropagationEvent({ callContext, tagRanges, result }) {
  const {
    _state: {
      orig: { args, obj },
      trackData
    }
  } = callContext;
  const signature = new Signature(PROPAGATOR_METHOD);
  const context = CallContext.create({
    obj,
    args,
    result,
    stackOpts: {
      constructorOpt: callContext.hooked
    }
  });
  const source = getSourceType(trackData);
  const parents = getParentEvents(source, trackData);
  const event = new PropagationEvent({
    context,
    parents,
    signature,
    source,
    tagRanges,
    target: 'R'
  });
  return event;
}

/**
 * Returns the parents for the event based on source.
 * @param {string} source The souce type
 * @param {object} trackData Tracking data
 * @returns {PropagationEvent[]|CreationEvent[]}
 */
function getParentEvents(source, trackData) {
  return _.compact([
    ...(source === 'P'
      ? [trackData.replacerData.event]
      : source === 'O'
        ? [trackData.objData.event]
        : [trackData.replacerData.event, trackData.objData.event]),
    ...Array.from(trackData.dynamicEvents)
  ]);
}

/**
 * Returns the source type for the event based on which consituents are tracked.
 * @param {object} callContext
 * @returns {string}
 */
function getSourceType({ objData, replacerData }) {
  return objData.tracked && !replacerData.tracked
    ? 'O'
    : !objData.tracked && replacerData.tracked
        ? 'P'
        : 'A';
}

module.exports.coerceReplacer = coerceReplacer;
module.exports.coerceEmptyReplacerArgs = coerceEmptyReplacerArgs;
module.exports.getMatchAndOffsetValues = getMatchAndOffsetValues;
module.exports.getTagRangesTagsList = getTagRangesTagsList;
module.exports.invokeReplacer = invokeReplacer;
