/**
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.
*/
/**
 * JSON.stringify Propagation Provider module.
 * @module lib/assess/propagators/JSON/stringify.js
 */
'use strict';

const crypto = require('crypto');

const TagRange = require('../../models/tag-range');
const tracker = require('../../../tracker.js');
const { isString } = require('../../../util/is-string');
const {
  PropagationEvent,
  Signature,
  CallContext,
  SourceEvent
} = require('../../models');
const ContrastJSON = require('../../../core/rewrite/injections').get('JSON');
const Debraner = require('../../membrane/debraner');
const { PROXY_TARGET } = require('../../../constants');
const { trackedObjects } = require('../../sources/event-handler');

const sig = new Signature({
  moduleName: 'JSON',
  methodName: 'stringify',
  isModule: false
});

const patcher = require('../../../hooks/patcher');
const { PATCH_TYPES } = require('../../../constants');

function makeCanary() {
  return crypto
    .randomBytes(12)
    .toString('base64')
    // map regex special chars to other chars
    .replace(/\+/g, '%')
    .replace(/\//g, '#');
}

// read canary as if the word was "marker". it's just a string inserted into the values
// so they can be located after stringifying the object.
const canary = makeCanary();

const regex = new RegExp(`"${canary}(\\d+)~`, 'g');

function getUntrustedSpaceProps(space) {
  // it can't be a problem if it's not a string, if the string is all spaces, or
  // if the string is zero length.
  // maybe match on a backslash or colon present? or maybe don't overthink - when
  // has user input ever been used for the json formatting padding string?
  if (!isString(space) || space.match(/^\s+$/) || !space) {
    return null;
  }
  // otherwise if the space string is tracked then the entire json output inherits
  // the tracked string's tags.
  const props = tracker.getData(space);
  if (!props || props.tagRanges.length === 0) {
    return null;
  }

  return props;
}

/**
 * Wrapper function to return the value when we patch the replacer to
 * always be a function
 *
 * @param {string} key
 * @param {any*} value
 * @return {any*} value
 */
function identity(key, val) {
  return val;
}

/**
 * Turns the replacer into a function so we can properly wrap the replacer later
 * w/ contrastReplacer
 *
 * @param {any*} replacer
 * @return {Function} wrapped or original replacer function
 */
function patchReplacer(replacer) {
  if (Array.isArray(replacer)) {
    // keys that are numbers get cast to string when they come through the replacer
    // add '' to the list of keys as that's always the first iteration of replacer, the entire json string
    const filterKeys = replacer
      .filter((key) => isString(key) || typeof key === 'number')
      .concat('')
      .map((value) => value.toString());

    let nextArray = [];
    replacer = function(key, value) {
      // we need to keep track of arrays
      // and shift each element as the replacer hits it
      // so we can compare those values are expected
      // when filtering
      const nextEl = nextArray.shift();
      if (Array.isArray(value)) {
        nextArray = [...value];
      }

      if (filterKeys.includes(key) || value === nextEl) {
        return value;
      }
    };
  } else if (typeof replacer === 'function') {
    // do nothing!
  } else {
    replacer = identity;
  }

  return replacer;
}

let patched = false;

module.exports.handle = function() {
  if (patched) {
    return;
  }
  patched = true;
  ContrastJSON.enable();
  patcher.patch(ContrastJSON.value, 'stringify', {
    name: ContrastJSON.name,
    patchType: PATCH_TYPES.ASSESS_PROPAGATOR,

    pre(data) {
      const [obj, , space] = data.args;
      let [input, replacer] = data.args;
      replacer = patchReplacer(replacer);

      // context used by the post hook.
      data.metadata = {
        origArgs: data.args.slice(),
        propagate: false,
        regex,
        spaceProps: undefined,
        target: undefined,
        membrane: undefined,
        debraner: undefined,
        sourcesMap: {}
      };

      if (trackedObjects.has(obj)) {
        const { type } = trackedObjects.get(obj);
        data.metadata.sourcesMap[type] = obj;
      }

      // is the argument behind a membrane? try getting { target, membrane }.
      const braned = input && input[PROXY_TARGET];

      const isSourceObject = input && trackedObjects.has(input);

      if (isSourceObject) {
        data.metadata.propagate = true;
      }

      // next two are used by the replacer function
      let safeKeys;
      if (braned) {
        data.metadata.propagate = true;
        data.metadata.braned = braned;
        ({ safeKeys } = braned.membrane);
        const debraner = new Debraner(input);
        data.metadata.debraner = debraner;
        // remove any membranes present in the input.
        input = debraner.remove();
      }

      // if the space can't be trusted then remember space's tracking info so it
      // can be applied to the entire JSON.stringify result string.
      const spaceProps = getUntrustedSpaceProps(space);
      if (spaceProps) {
        data.metadata.spaceProps = spaceProps;
        data.metadata.propagate = true;
      }

      // each canary tag gets a separate id to make matching easier.
      let id = 0;

      /**
       * wraps the replacer method and updates propagate boolean if value of replacer
       * is tracked
       *
       * @param {string} key
       * @param {string} value
       */
      // eslint-disable-next-line complexity
      function contrastReplacer(key, val) {
        if (trackedObjects.has(obj)) {
          const { type } = trackedObjects.get(obj);
          if (!data.metadata.sourcesMap[type]) {
            data.metadata.sourcesMap[type] = obj;
          }
        }

        let isTracked = false;
        const valProperties = tracker.getData(val);
        if (valProperties && valProperties.tagRanges.length) {
          data.metadata.propagate = true;
          isTracked = true;
        }

        let value = replacer.call(this, key, val);
        if (!isTracked) {
          if (!safeKeys) {
            return value;
          }
          const safeProps = safeKeys.get(this);
          if (!safeProps || !safeProps.has(key)) {
            return value;
          }
          // it's not tracked but it is in safeProps, so it should be trusted. an
          // empty tag ('') is a flag indicating that the string is trusted.
          if (isString(value)) {
            // make skeletal version of valProperties with only the tagRanges prop.
            data.metadata[id] =
              { tagRanges: [new TagRange(0, value.length - 1, '')] };
            value = `${canary}${id}~${value}`;
            id++;
          }
          return value;
        }

        // keys aren't propagated, so check only the value
        if (isString(value)) {
          data.metadata[id] = valProperties;
          value = `${canary}${id}~${value}`;
          id++;
        }
        return value;
      }

      data.args = [input, contrastReplacer, space];
    },

    // eslint-disable-next-line complexity
    post(data) {
      if (!data.metadata.propagate) {
        return data.result;
      }

      // restore the proxies to the argument
      const { debraner } = data.metadata;
      if (debraner) {
        debraner.restore();
      }

      const { metadata } = data;
      // tag ranges will be set to the spaceProps properties if present,
      // else each property value's tags as recorded by the replacer.
      let tagRanges = [];

      // if the space argument was an untrusted string then its tagRanges
      // apply to the entire result. we just invalidate the entire output
      // string because we can't know what was impacted without doing deep
      // analysis.
      if (metadata.spaceProps) {
        data.result = applySpaceTagRanges(data.result, metadata, tagRanges);
      } else {
        data.result = applyCanaryTagRanges(data.result, metadata, tagRanges);

        // now, if the object being stringified was behind a membrane, add
        // the membrane's tag(s) around everything already tagged.
        if (debraner) {
          tagRanges = weaveMembraneTagRanges(tagRanges, data.result.length);
        }
      }

      const tracked = tracker.track(data.result);
      if (!tracked) {
        return data.result;
      }
      data.result = tracked.str;
      const { props } = tracked;

      const customSources = [];

      Object.keys(data.metadata.sourcesMap).forEach((k) => {
        const objData = trackedObjects.get(data.metadata.sourcesMap[k]);

        const { type, path, stackOpts } = objData;

        const obj = data.metadata.sourcesMap[type];

        customSources.push(
          new SourceEvent({
            context: CallContext.create({
              obj,
              args: data.args,
              result: data.result,
              stackOpts
            }),
            signature: new Signature({
              moduleName: 'http.req',
              methodName: 'on.end',
              args: undefined,
              return: undefined,
              isModule: false
            }),
            tagRanges,
            source: 'A',
            target: 'R',
            name: path,
            type,
            isRequestObject: true
          })
        );
      });

      props.tagRanges = tagRanges;
      // restore the original arguments so reporting shows user's signature.
      data.args = metadata.origArgs;

      // this check is required only for tests, i believe. sourceEvents should
      // always exist, even if an empty array, for production code.
      if (metadata.sourceEvents && !metadata.sourceEvents.length) {
        const { sourceEvents, braned } = data.metadata;
        // If it is a source membrane (could be deserialization membrane), then
        // we need to get a source event.
        if (braned && braned.membrane && braned.membrane.makeReqSourceEvent) {
          sourceEvents.push(
            braned.membrane.makeReqSourceEvent(data.result.length)
          );
        }
      }

      props.event = new PropagationEvent({
        context: CallContext.create(data),
        signature: sig,
        tagRanges: props.tagRanges,
        source: 'P',
        target: 'R',
        parents: customSources.length
          ? customSources
          : data.metadata.sourceEvents
      });

      return data.result;
    }
  });

  function applySpaceTagRanges(result, metadata, tagRanges) {
    // remove the markers that were inserted. this could be optimized by not
    // inserting the markers to begin with, but this is a degenerate case to
    // begin with.
    result = result.replace(metadata.regex, (m) => '"');
    for (const tagRange of metadata.spaceProps.tagRanges) {
      tagRanges.push(new TagRange(0, result.length - 1, tagRange.tag));
    }
    return result;
  }

  function applyCanaryTagRanges(result, metadata, tagRanges) {
    metadata.sourceEvents = [];
    // the diff doesn't include the " that opens the regex.
    let canaryReplacementDiff = 0;
    // for each marker in the stringify's result, remove the marker and
    // create a TagRange for the value.
    return result.replace(metadata.regex, (m, id, offset) => {
      // adjust offset by total characters removed so far
      offset = offset - canaryReplacementDiff + 1;
      // we don't replace the opening " in the regex - that just makes sure
      // we only find our marker at the start of a string value.
      canaryReplacementDiff += m.length - 1;
      for (const tr of metadata[id].tagRanges) {
        tagRanges.push(new TagRange(tr.start + offset, tr.stop + offset, tr.tag));
      }

      // no event will be present if the tag range is the null tag (''). that is
      // used for trusted strings that have been added to the object behind the
      // membrane. there is no event associated with them; they're a flag for
      // weaveMembraneTagRanges() to skip those ranges.
      if (metadata[id].event) {
        metadata.sourceEvents.push(metadata[id].event);
      }
      // replace the canary with the starting quote
      return '"';
    });
  }

  /**
   * function to weave tag ranges found on properties of the stringified
   * object into an implicit global "untrusted" tag range covering every
   * other part of the stringified object.
   */
  function weaveMembraneTagRanges(tagRanges, resultLength) {
    // find any tags that don't match untrusted
    const otherTags = tagRanges.filter((tr) => tr.tag !== 'untrusted');

    // if there are no tags other than untrusted, no need to weave them in.
    if (!otherTags.length) {
      return [new TagRange(0, resultLength - 1, 'untrusted')];
    }

    const woven = [];
    // sorting by start tag should be enough as there cannot be overlap.
    otherTags.sort((a, b) => a.start - b.start);
    let ix = 0;
    while (otherTags.length) {
      const tag = otherTags.shift();
      if (tag.start > ix) {
        const stop = tag.start - 1;
        woven.push(new TagRange(ix, stop, 'untrusted'));
        ix = tag.stop + 1;
        // an empty tag means no tag, so don't add it.
        if (tag.tag) {
          woven.push(tag);
        }
      } else if (tag.start === ix) {
        // this shouldn't happen, i don't think. there should
        // always be a double-quote, comma, double-quote...
        ix = tag.stop + 1;
        if (tag.tag) {
          woven.push(tag);
        }
      } else {
        // this definitely shouldn't happen
      }
    }
    // catch any trailing untrusted.
    if (ix < resultLength) {
      woven.push(new TagRange(ix, resultLength - 1, 'untrusted'));
    }

    return woven;
  }
};
