/**
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 querystring = require('querystring');
const _ = require('lodash');
const tracker = require('../../../tracker');
const TagRange = require('../../models/tag-range');
const Scopes = require('../../../core/async-storage/scopes');
const { CallContext, PropagationEvent, Signature } = require('../../models');
const tagRangeUtil = require('../../models/tag-range/util');

/**
 * Helper class to manage state which responds to calls to the escape function.
 * We know that escape will get called once per key and once per value, and that
 * keys are escaped first. So then, based on number of calls we can determine
 * the position in the result string which corresponds to a value being escaped.
 *
 * Example:
 *   `querystring.stringify({ x: 'foo' y: ['bar', 'baz', 'qux'] })`
 *
 * will be processed
 *
 *        [0]  x=
 *        [2]  . foo&
 *        [6]  . .   y=
 *        [8]  . .   . bar&
 *       [14]  . .   . .     baz&
 *       [20]  . .   . .     .     qux
 *             . .   . .     .     .
 *             x=foo&y=bar&y=baz&y=qux
 *             | |   | |     |     |
 *  [offsets]  0123456789012345678901
 *             0         1         2
 *
 * Since we have this offset, we can simply return TagRanges for the escaped
 * component with the offset and adjust them as the result is being processed.
 * This approach relies on running the escape function in an ALLOW_ALL scope,
 * otherwise propagation in it would be thwarted by current PROPAGATION scope.
 */
class OnEscapeHandler {
  /**
   * @param {Object} input input argument to stringify
   * @param {string} equals value used as equals (=)
   * @param {string} separator value used as separator (&)
   */
  constructor({ input, equals, separator }) {
    this.input = input;
    this.equals = equals;
    this.separator = separator;
    this.currentKey = { tagRanges: null, value: null };
    this.readingMultipleValues = false;
    this.valuesRemaining = 0;
    this.offset = 0;
    this.remainingCalls = OnEscapeHandler.getExpectedCallCount(input);
  }

  /**
   * Shorthand to get the equals string value's length.
   * @returns {number}
   */
  get eqLen() {
    return this.equals.length;
  }

  /**
   * Shorthand to get the separator string value's length. This value is
   * state-dependent.
   * @returns {number}
   */
  get sepLen() {
    return this.remainingCalls === 0 ? 0 : this.separator.length;
  }

  /**
   * When there are no more values for the current key, we must be processing a
   * new key.
   * @param {string} component
   * @returns {boolean}
   */
  isKey(raw) {
    return this.valuesRemaining === 0 && raw in this.input;
  }

  /**
   * Called each time encodeURIComponent is called. Using state management, we
   * will process the information as a key or a value and return TagRanges that
   * will need to be transferred.
   * @param {string} opts.raw
   * @param {string} opts.escaped
   * @returns {Object[]}
   */
  getTrackingData({ raw, escaped }) {
    let ret;
    this.remainingCalls--;
    if (this.isKey(raw)) {
      ret = this.handleKey({ raw, escaped });
    } else {
      ret = this.handleValue({ raw, escaped });
    }
    return ret;
  }

  /**
   * Logic for handling when escape is called on object key. A call to this
   * corresponds to querystring.stringify creating the first element of a
   * tuple, such as `key_n=`.
   * @param {string} raw The raw key value
   * @param {string} escaped Value returned from escape function
   * @returns {Object[]}
   */
  handleKey({ raw, escaped }) {
    let ret;
    const value = this.input[raw];
    const trackingData = tracker.getData(escaped);

    // set current key for reference if when array-valued
    this.currentKey = {
      value: escaped,
      tagRanges: trackingData ? trackingData.tagRanges : null
    };

    // capture track info before updating state
    if (this.currentKey.tagRanges) {
      ret = [
        {
          offset: this.offset,
          tagRanges: trackingData.tagRanges,
          event: trackingData.event
        }
      ];
    }

    // update offset and value-tracking state
    this.offset += escaped.length + this.eqLen;
    this.valuesRemaining = Array.isArray(value) ? value.length : 1;
    this.readingMultipleValues = false;

    return ret || [];
  }

  /**
   * Logic for handling when escape is called on object key value. A call to
   * this corresponds to querystring.stringify creating the second element of a
   * tuple with the separator: `value&`.
   * @param {string} raw the raw value
   * @param {string} escaped
   * @returns {Object[]}
   */
  handleValue({ raw, escaped }) {
    const trackingData = tracker.getData(escaped);
    const ret = [];

    if (this.readingMultipleValues) {
      // capture ranges for subsequent appearances of keys when they are array-valued.
      // do that first and then adjust the offset.
      if (this.currentKey.tagRanges) {
        ret.push({
          offset: this.offset,
          tagRanges: this.currentKey.tagRanges
        });
      }
      this.offset += this.currentKey.value.length + this.eqLen;
    }

    if (trackingData) {
      ret.push({
        offset: this.offset,
        tagRanges: trackingData.tagRanges,
        event: trackingData.event
      });
    }

    // update state
    this.offset += escaped.length + this.sepLen;
    this.valuesRemaining--;
    this.readingMultipleValues = true;

    return ret;
  }

  /**
   * How many times will escapeURIComponent be called for a given input.
   * @param {Object} obj
   */
  static getExpectedCallCount(input) {
    let valueCount = 0;
    const keys = Object.keys(input);
    keys.forEach((key) => {
      if (Array.isArray(input[key])) {
        valueCount += input[key].length;
      } else {
        valueCount += 1;
      }
    });
    return keys.length + valueCount;
  }
}

/**
 * Sets up propagator state and weaves an instrumentation function into the
 * propagator's arguments as the custom encoding function. The instrumentation
 * will collect events and offset-adjusted TagRanges with each call to the
 * escape function.
 * @param {Object} data
 */
function pre(data) {
  const input = data.args[0];

  if (!input) {
    return;
  }

  const separator = data.args[1] || '&';
  const equals = data.args[2] || '=';
  const handler = new OnEscapeHandler({ input, separator, equals });
  const escape =
    data.args[3] && typeof data.args[3].encodeURIComponent === 'function'
      ? data.args[3].encodeURIComponent
      : querystring.escape;

  // setup state
  data.state = {
    escape,
    // because we alter them
    origArgs: [...data.args],
    // collect parent events
    events: new Set(),
    // collect tag ranges
    tagRanges: []
  };

  // force the custom escape fn option as instrumentation mechanism
  data.args[3] = {
    ...(data.args[3] || {}),
    encodeURIComponent(...args) {
      const escaped = Scopes.runInAllowAllScope(() => escape(...args));
      const trackingData = handler.getTrackingData({ raw: args[0], escaped });

      // each call to the escape function may result in tags at diff offsets
      if (trackingData.length) {
        for (const { offset, tagRanges, event } of trackingData) {
          data.state.events.add(event);
          tagRangeUtil.addAllWithOffsetInPlace(
            data.state.tagRanges,
            tagRanges,
            offset
          );
        }
      }

      return escaped;
    }
  };
}

/**
 * Merges all collected TagRanges and collects events. Builds the new
 * propagation event and tracks the result.
 * @param {Object} data
 */
function post(data) {
  const input = data.args[0];

  if (!input || !data.result) {
    return;
  }

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

  if (data.state.escape === querystring.escape) {
    data.state.tagRanges.push(
      new TagRange(0, data.result.length - 1, 'url-encoded')
    );
  }
  if (tracked) {
    const sorted = _.sortBy(data.state.tagRanges, 'start');
    tracked.props.tagRanges = [];
    tagRangeUtil.addAllInPlace(tracked.props.tagRanges, sorted);

    // stringify / encode
    const method = data.funcKey.split('.')[1];

    tracked.props.event = new PropagationEvent({
      context: CallContext.create({
        ...data,
        args: data.state.origArgs,
        obj: null
      }),
      parents: Array.from(data.state.events),
      signature: new Signature(`querystring.${method}`),
      source: 'P',
      tagRanges: tracked.props.tagRanges,
      target: 'R',
      tags: ['url-encoded']
    });
    data.result = tracked.str;
  }
}

module.exports.handle = { pre, post };
