/**
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 _ = require('lodash');
const requireHook = require('../../../hooks/require');
const patcher = require('../../../hooks/patcher');
const {
  PATCH_TYPES: { ASSESS_PROPAGATOR },
} = require('../../../constants');
const tracker = require('../../../tracker');
const tagRangeUtil = require('../../models/tag-range/util');
const { toUntrackedString } = require('../../utils');
const logger = require('../../../core/logger')(
  'contrast:assess:propagators:joi'
);

/**
 * Input strings validated through schemas having `.allow()`, `.valid()`, and
 * `.equal()` should not be tracked. The implementation in joi around this
 * feature involves the `Values` class which is an abstraction on `Set`.
 * If the value being validated is not found when `values.get(v)`is called it
 * returns `false`. When it is found will return `{ value: v }`. Any time the
 * return object's `value` property is a tracked string we can untrack it.
 * @param {Values} values
 */
function instrumentJoiValues(values) {
  patcher.patch(values.prototype, 'get', {
    name: 'joi.values',
    patchType: ASSESS_PROPAGATOR,
    post(data) {
      const {
        args: [value],
        result,
      } = data;

      // value not found during lookup
      if (!result) {
        return;
      }

      if (result.ref) {
        const resultValue = result.ref.__CONTRAST__
          ? result.ref.__CONTRAST__.proxyTaggedString
          : result.value;

        handler(resultValue, value, data);

        if (result.ref.__CONTRAST__) {
          delete result.ref.__CONTRAST__;
        }

      } else if (_.isString(result.value)) {
        // use case is .valid() - safe
        result.value = tracker.untrack(result.value) || result.value;
      }
    },
  });
}

const stringHandler = (resultValue, argValue, data) => {
  const resultIsString = _.isString(resultValue);
  const argIsString = _.isString(argValue);

  if (resultIsString && argIsString) {
    const resolvedTrackData = tracker.getData(resultValue);
    const refTrackData = tracker.getData(argValue);
    const handler = getRefHandler(resolvedTrackData, refTrackData);
    handler && handler(data, resolvedTrackData, refTrackData);
  }
};

const handler = (resultValue, argValue, data) => {
  if (_.isString(resultValue)) {
    return stringHandler(resultValue, argValue, data);
  }

  if (_.isObject(resultValue)) {
    for (const [key, value] of Object.entries(resultValue)) {
      handler(value, argValue[key], data);
    }
  }
};

/**
 * Depending on which values are tracked, ref and/or target, returns the
 * appropriate function to handle the scenario.
 * @param {object} resolvedTrackData tracking data of resolved target
 * @param {object} refTrackData tracking data of ref
 * @returns {function}
 */
function getRefHandler(resolvedTrackData, refTrackData) {
  // 4 Cases
  if (!resolvedTrackData) {
    if (!refTrackData) {
      return null;
    } else {
      return handleRefOnlyTracked;
    }
  } else {
    if (refTrackData) {
      return handleBothTracked;
    } else {
      return handleTargetOnlyTracked;
    }
  }
}

/**
 * Prevents natural tendency for untracked reference to become tracked in
 * validation result.
 * @param {object} data patcher wrap context
 * @param {object} resolvedTrackData tracking data for resolved value
 * @param {object} refTrackData tracking data for reference value
 */
function handleTargetOnlyTracked(data, resolvedTrackData, refTrackData) {
  const {
    args: [value],
  } = data;
  data.result.value = toUntrackedString(value);
}

/**
 * Apply the same validation events to the ref that were applied to the resolved
 * value.
 * @param {object} data patcher wrap context
 * @param {object} resolvedTrackData tracking data for resolved value
 * @param {object} refTrackData tracking data for reference value
 */
function handleBothTracked(data, resolvedTrackData, refTrackData) {
  let value = data.args[0];
  const {
    args: [, , prefs],
    result,
  } = data;

  // We can't reliably validate values that get adjusted
  if (result.ref && result.ref.adjust) {
    return;
  }

  if (result.ref.map) {
    data.result.value = tracker.untrack(data.result.value) || data.result.value;
    value = tracker.untrack(value) || value;
  } else {
    copyValidationHistory(resolvedTrackData, refTrackData);
    if (prefs.convert) {
      // this prevents the ref from being set to the resolved target value
      data.result.value = value;
    }
  }
}

/**
 * Handles when only the ref()'s string is tracked.
 * @param {object} data patcher wrap context
 * @param {object} resolvedTrackData tracking data for resolved value
 * @param {object} refTrackData tracking data for reference value
 */
function handleRefOnlyTracked(data, resolvedTrackData, refTrackData) {
  let value = data.args[0];
  const {
    args: [, , prefs],
    result,
  } = data;

  if (prefs.convert) {
    // If there's no mapping, always just prevent ref's value from  being
    // replaced as to maintain ref tracking
    if (!result.ref.map) {
      data.result.value = value;
      // TODO(NODE-1504): How to determine validation tags in this case?
      logger.debug(
        'unable to validate ref of untracked target: %s',
        result.ref.display
      );
    }
  } else {
    // if map is used we can trust - like .valid()
    if (result.ref.map) {
      if (!tracker.getData(result.value)) {
        value = tracker.untrack(value) || value;
      }
    } else {
      logger.debug(
        'unable to validate ref of untracked target: %s',
        result.ref.display
      );
    }
  }
}

/**
 * References should have the same Joi validation tags/events applied as their
 * resolved targets. We can trust that the target has had all joi validation
 * rules applied before the reference is resolved. The basic idea is:
 * 1. Starting at the target's last event, get all of the contiguous joi
 *    validation events.
 * 2. From the events determine which events applied which tags.
 * 3. Update the ref's tracking data with those events and tag ranges - as if
 *    they were "replayed".
 * @param {object} targetTrackData tracker data for ref target
 * @param {object} refTrackData tracker data for ref
 */
function copyValidationHistory(targetTrackData, refTrackData) {
  const joiEvents = getContiguousJoiEvents(targetTrackData.event);
  const validationHistory = buildEventsAndTagsToReplay(joiEvents);
  applyValidationHistoryToRef(refTrackData, validationHistory);
}

/**
 * Given a list of events will attempt to determine which tag ranges were added
 * as a result of each event occuring. Events have a `tagRanges` property but
 * they contain the entire set of tag ranges for the string after the event
 * occured. So we have to do some work to figure out which tag ranges were added
 * by each event before we can "replay" them on the ref tracking data. Then we
 * return info in the form [{ event, tagRanges }, ...].
 * @param {object[]} joiEvents the `PropagationEvent`'s in target's history
 * @returns {object[]}
 */
function buildEventsAndTagsToReplay(joiEvents) {
  const hist = [];
  const tagRangesSeen = new WeakSet(_.get(joiEvents, '0.parents.0.tagRanges'));

  for (const event of joiEvents) {
    const newEventRanges = [];

    // collect any new tagRanges added for event
    for (const rCurr of event.tagRanges) {
      if (tagRangesSeen.has(rCurr)) {
        continue;
      }
      tagRangesSeen.add(rCurr);
      newEventRanges.push(rCurr);
    }

    // build into history
    if (newEventRanges.length) {
      hist.push({
        event,
        tagRanges: newEventRanges,
      });
    }
  }

  return hist;
}

/**
 * Returns the set of contiguous joi propagation events starting at the provided
 * event.
 * @param {Event} event a dataflow event (source|propagation)
 * @param {array} coll collection built during recursion
 * @returns {array}
 */
function getContiguousJoiEvents(event, coll = []) {
  if (!event || !event.signature) {
    return coll;
  }
  if (event.signature.moduleName === 'joi') {
    coll.unshift(event);
    return getContiguousJoiEvents(event.parents[0], coll);
  } else {
    return coll;
  }
}

/**
 * Given the ref's tracking data and the target's joi validation history, will
 * transfer the history details to the ref tracking data. The ref never
 * undergoes validation itself but rather by proxy.
 * @param {obj} refTrackData the ref's tracking data
 * @param {object[]} hist array of contiguous Joi validation events/tagRanges
 */
function applyValidationHistoryToRef(refTrackData, hist = []) {
  for (const { event, tagRanges } of hist) {
    const { event: refParentEvent } = refTrackData;
    refTrackData.event = _.cloneDeep(event);
    refTrackData.event.parents = [refParentEvent];
    for (const tagRange of tagRanges) {
      refTrackData.tagRanges = tagRangeUtil.add(
        refTrackData.tagRanges,
        tagRange
      );
    }
  }
}

requireHook.resolve(
  { name: 'joi', file: 'lib/values.js', version: '>=17.0.0' },
  instrumentJoiValues
);
