/**
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 logger = require('../../core/logger')('contrast:source-membrane');
const Membrane = require('./index');
const tracker = require('../../tracker');
const {
  SourceEvent,
  PropagationEvent,
  CallContext,
  Signature
} = require('../models');
const { AsyncStorage, KEYS } = require('../../core/async-storage');
const Scopes = require('../../core/async-storage/scopes');
const TraceEventSource = require('../../reporter/models/trace-event-source');
const TagRange = require('../models/tag-range');
const patcher = require('../../hooks/patcher');
const { PATCH_TYPES } = require('../../constants');

const signature = new Signature({
  moduleName: 'Object',
  methodName: 'getter',
  args: ['Object'],
  return: 'String',
  isModule: false
});

module.exports = class SourceMembrane extends Membrane {
  /**
   * @param {object} config agent config to use for array sampling.
   * @param {object} options
   * @param {object} options.inputExclusions active input exclusions for membrane type e.g. PARAMETER
   * @param {string} options.type part of request object: 'BODY', 'HEADER', etc.
   * @param {function} options.constructorOpts argument to limit stack trace reporting (see stacktrace.js)
   */
  constructor(conf, options) {
    const cfg = conf.agent.node;
    super(
      'enable_property_descriptor_tracking' in cfg
        ? cfg.enable_property_descriptor_tracking
        : true
    );
    // because we don't actually 'wrap' arrays (if sampling) so
    // we need another way to prevent double wrapping,
    // or doing extra work on consecutive array references

    this.wrappedArrays = new WeakSet();
    this.inputExclusions = options.inputExclusions || [];
    this.deadzoneEnabled = cfg.enable_sources_in_deadzones;
    if (cfg.array_request_sampling) {
      this.sample = cfg.array_request_sampling.enable;
      this.sampleThreshold = cfg.array_request_sampling.threshold;
      this.sampleInterval = cfg.array_request_sampling.interval || 1;
    }
    // used when a reqSourceEvent must be created because an entire object
    // is being stringified but individual properties are not being referenced
    // due to debraning.
    this.type = options.type.toLowerCase();
    const { object, constructorOpts } = options;
    this.callContext = this.makeCallContext(object, constructorOpts);
  }

  /**
   * Derived from the metadata, store source event so that it can be accessed at
   * route report time.
   * spec: https://bitbucket.org/contrastsecurity/assess-specifications/src/master/remediation/route-remediation.md
   * @param {object} metadata source event data to be stored
   */
  static storeSource(metadata) {
    const { sourceType: type, path: name } = metadata;
    let storedSources = AsyncStorage.get(KEYS.SOURCES);

    if (!storedSources) {
      // if this is the first source stored on the request, initialize a map
      // and set the storedSources to be a reference to the new, empty map.
      AsyncStorage.set(KEYS.SOURCES, new Map());
      storedSources = AsyncStorage.get(KEYS.SOURCES);
    }

    // there are cases where AsyncStorage context is out of sync, add defensive code so we don't
    // crash.
    if (storedSources) {
      // save event for route coverage (data pulled in for RouteInfo object)
      // only want to save unique values
      storedSources.set(
        `${type}-${name}`,
        new TraceEventSource({ type, name })
      );
    }
  }

  /**
   * Used to wrap a buffer's toString to maintain taint
   * when a buffer is a property of a tainted object
   *
   * @param  {Buffer} buf      buffer we're tracking
   * @param  {Object} metadata sourceType and path in object
   * @return {Buffer}          patched instance
   */
  wrapBuffer(buf, metadata = {}) {
    patcher.patch(buf, 'toString', {
      name: 'Buffer.toString',
      alwaysRun: true,
      patchType: PATCH_TYPES.MISC,
      post: (wrapCtx) => {
        wrapCtx.result = this.onString(wrapCtx.result, metadata);
      }
    });

    return buf;
  }

  /**
   * Used to create a SourceEvent with relevant TagRanges when a key on an
   * object is referenced. We cache tracking data of values accessed in order
   * to be able to preserve tag ranges and event history if P->A/O events
   * propagate "upwards" to source event.
   * @param {string} str value of key on object getting referenced
   * @param {object} metadata sourceType and path in object
   * @return {string} returns tracked string with proper contrastProps
   */
  onString(str, metadata) {
    if (!this.ensureMetadata(metadata) && !metadata.rootLevelString) {
      return str;
    }
    const tracked = tracker.track(str);
    if (!tracked) {
      return str;
    }
    this.setTrackingData(tracked, metadata);
    SourceMembrane.storeSource(metadata);

    return tracked.str;
  }

  makeCallContext(object, constructorOpt) {
    // capture the source of all the properties now, not when a string is
    // referenced.
    const context = CallContext.create({
      obj: 'IncomingMessage',
      args: [''],
      result: object,
      stackOpts: {
        constructorOpt
      }
    });
    // set manually as this is executed before wrapping in a membrane.
    context.resultTracked = true;

    return context;
  }

  makeReqSourceEvent(length) {
    const signature = new Signature({
      moduleName: 'http.req',
      methodName: 'on.end',
      args: undefined,
      return: undefined,
      isModule: false
    });

    const reqSourceEvent = new SourceEvent({
      context: this.callContext,
      signature,
      tagRanges: [{ start: 0, stop: length - 1, tag: 'untrusted' }],
      source: 'A',
      target: 'R',
      type: this.type,
      name: '',
      isRequestObject: true
    });

    return reqSourceEvent;
  }

  /**
   * Given a tracked string will check cache for any events/tag ranges which have been
   * cached from previous key accesses. If no cached data is found will construct new
   * source event/tag ranges and apply to tracking data, then cache them to be reused for
   * next time key is accessed. If the first propagation event that a membrane's value
   * should go through adds tags/event to the source value (e.g. a validator such as
   * joi.string().validate(req.query.input)), then the cached tracking data would have an
   * event that is of type PropagationEvent. In these cases we need to copy the full event
   * history and replace the SourceEvent in that history with a new copy. This will allow
   * the value being accessed to have the correct stacktrace without effecting the stacktrace
   * of the same value being accessed from the membrane from a different code location.
   * @param {string} tracked tracked string being accessed
   * @param {object} metadata sourceType, last target, dot-delimitted path of value being accessed
   */
  setTrackingData(trackResult, metadata) {
    const { str: tracked, props: trackData } = trackResult;
    const context = CallContext.create({
      obj: 'IncomingMessage',
      args: [metadata.path],
      result: tracked,
      stackOpts: {
        constructorOpt: this.wrap
      }
    });
    const cached = this.getCachedTrackingData(metadata);
    if (cached) {
      trackData.event = this.copyCachedEventHistory(cached.event, context);
      trackData.tagRanges = Array.from(cached.tagRanges);
    } else {
      trackData.tagRanges = this.getTagRanges(tracked, metadata);
      trackData.event = new SourceEvent({
        context,
        signature,
        tagRanges: trackData.tagRanges,
        target: 'R',
        type: metadata.sourceType,
        name: metadata.path,
        isRequestObject: false,
        isArray: metadata.isArray
      });
    }
    // update cache with current tracking data for returned string
    this.cacheTrackingData(trackData, metadata);
  }

  /**
   * Creates a copy of the tracking data's event history for the string being
   * accessed from the membrane. At this point in the dataflow the history must
   * be linear, so we only traverse up the existing history by 1st parent event.
   * @param {SourceEvent|PropagationEvent} cachedEvent cached SourceEvent or PropagationEvent
   * @param {CallContext} context
   * @returns {SourceEvent|PropagationEvent}
   */
  copyCachedEventHistory(cachedEvent, context) {
    if (cachedEvent instanceof SourceEvent) {
      return new SourceEvent({ ...cachedEvent, context });
    }

    // we need to return this
    const parentCopy = new PropagationEvent(cachedEvent);

    // state for traversal
    let nextCopy = parentCopy;
    let prevCopy;
    let nextOrigEvent = cachedEvent;
    let prevOrigEvent;

    while (nextOrigEvent && nextOrigEvent.parents.length) {
      prevOrigEvent = nextOrigEvent;
      nextOrigEvent = prevOrigEvent.parents[0];
      if (nextOrigEvent instanceof SourceEvent) {
        // replace the last copied PropagationEvent's parent with a copy of the
        // source event having new context
        nextCopy.parents = [new SourceEvent({ ...nextOrigEvent, context })];
        nextOrigEvent = null;
      } /* PropagationEvent */ else {
        prevCopy = nextCopy;
        nextCopy = new PropagationEvent(nextOrigEvent);
        prevCopy.parents = [nextCopy];
      }
    }
    return parentCopy;
  }

  /**
   * The way we wrap query params in koa is a little funky
   * @link https://github.com/koajs/koa/blob/master/lib/request.js#L172
   * Koa wraps query params in a cache so if a request has a query param of
   * `input=foo`
   * the path will be `input=foo.input` we want to fix that for call context to just be
   * `input`
   * @param {} metadata
   */
  ensureMetadata(metadata) {
    if (!(metadata.sourceType && metadata.path)) {
      return false;
    }
    const koaQueryString = metadata.path.split('=');
    if (koaQueryString[1]) {
      metadata.path = koaQueryString[0];
    }
    return true;
  }

  /**
   * The `TagRanges` returned  will include the `untrusted` tag. There will be
   * `exclusion:${ruleId}` tags for input exclusions pertaining to specific
   * rules and whose name matches the property name described in metadata.
   * @param {string} str string being tracked by membrane
   * @param {object} metadata metadata about source type and key name
   * @returns {TagRange[]}
   */
  getTagRanges(str, metadata) {
    const start = 0;
    const stop = str.length - 1;
    const tagRanges = [new TagRange(start, stop, 'untrusted')];

    if (this.inputExclusions.length) {
      const rules = new Set();
      const exclusions = this.inputExclusions.filter(
        (e) => !e.appliesToAllAssessRules()
      );

      for (const exclusion of exclusions) {
        if (exclusion.matches(metadata.path)) {
          for (const ruleId of exclusion.assessmentRulesList) {
            logger.debug(
              'excluding %s %s value from %s (%s)',
              metadata.sourceType,
              metadata.path,
              ruleId,
              exclusion.name
            );

            if (!rules.has(ruleId)) {
              tagRanges.push(new TagRange(start, stop, `exclusion:${ruleId}`));
            }

            rules.add(ruleId);
          }
        }
      }
    }

    if (metadata.sourceType === 'header') {
      if (metadata.path.toLocaleLowerCase() !== 'referer') {
        tagRanges.push(new TagRange(start, stop, 'header'));
      }
    } else if (metadata.sourceType === 'cookie') {
      tagRanges.push(new TagRange(start, stop, 'cookie'));
    }

    return tagRanges;
  }

  /**
   * Used to apply sampling on a membrane that is an array.
   * It wraps a sampling of values of array
   *
   * @param {Array} array value of key on object getting referenced
   * @param {Object} metadata sourceType and path in object
   * @return {Array} returns original array
   */
  wrapArray(arr, metadata) {
    metadata.isArray = true;

    // don't sample more than once
    if (this.wrappedArrays.has(arr)) {
      return arr;
    }

    const limit = Math.min(this.sampleThreshold, arr.length);

    // if not sampling, treat it as any object. not sampling has 100% accuracy.
    if (!this.sample || (limit === arr.length && this.sampleInterval === 1)) {
      return super.wrapObject(arr, metadata);
    } else if (!this.sampleThreshold) {
      return arr;
    } else {
      const origPath = metadata.path;
      for (let i = 0; i < limit; i += this.sampleInterval) {
        if (arr[i]) {
          const m = Object.assign({}, metadata);
          m.path = origPath ? `${origPath}[${i}]` : `[${i}]`;
          const wrapped = this.wrap(arr[i], m);
          arr[i] = wrapped;
        }
      }
    }

    this.wrappedArrays.add(arr);
    return arr;
  }

  /**
   * Short circuit the proxy and override the result
   *
   * @returns result non-null if overriding, null to continue
   */
  postResultOverride(prop, tar, result) {
    if (!this.deadzoneEnabled && Scopes.isNoInstrumentationScoped()) {
      return result;
    }
    // In the case where string properties are accessed and the membrane is
    // configured with exclusions, check the key to see if the value should be
    // excluded. If so we can skip tracking, but only if the exclusion is for
    // ALL rules. Otherwise, we'll account for the exclusion during onString by
    // adding exclusion tags for specific rules e.g. exclusion:<ruleId>.
    if (this.inputExclusions.length && typeof result === 'string') {
      const inputMatch = this.inputExclusions.find(
        (e) => e.appliesToAllAssessRules() && e.matches(prop)
      );

      if (inputMatch) {
        logger.debug(
          'excluding %s %s from all assess rules (%s)',
          inputMatch.inputType.toLowerCase(),
          prop,
          inputMatch.name
        );
        return result;
      }
    }

    return super.postResultOverride(prop, tar, result);
  }
};
