/**
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:assess.sources');
const { PATCH_TYPES } = require('../../constants');
const { CallContext, SourceEvent } = require('../models');
const TraceEventSource = require('../../reporter/models/trace-event-source');
const { AsyncStorage, KEYS } = require('../../core/async-storage');
const TagRange = require('../models/tag-range');
const patcher = require('../../hooks/patcher');
const tracker = require('../../tracker');
const parseurl = require('parseurl');

const SOURCE_EVENT_MAX = 250;
const trackedObjects = new WeakMap();

class SourceEventHandler {
  constructor({
    config = {
      agent: {
        node: {
          array_request_sampling: {
            enable: true,
            threshold: 50,
            interval: 5
          }
        }
      }
    },
    ensureReq = true,
    exclusions,
    stackOpts,
    snapshot,
    signature,
    type = 'UNKNOWN'
  } = {}) {
    this.config = config;
    this.exclusions = exclusions;
    this.type = type;
    this.snapshot = snapshot;
    this.signature = signature;
    this.stackOpts = stackOpts;
    this.ensureReq = ensureReq;
    this.samplingEnabled = config.agent.node.array_request_sampling.enable;
    this.samplingThreshold = config.agent.node.array_request_sampling.threshold;
    this.samplingInterval = this.samplingEnabled
      ? config.agent.node.array_request_sampling.interval
      : 1;

    this.inputExclusions = [];
    this.urlExclusions = [];
    this.skipEvent = false;
    this.sourceEventCount = 0;

    this.init();
  }

  init() {
    // values reused in functions
    this.req = AsyncStorage.get(KEYS.REQ);

    // NODE-1431: prevents crashing when req is not present in AsyncStorage.
    if (!this.req) {
      if (!this.ensureReq) return;

      this.skipEvent = true;
      logger.debug(
        'failed to handle source event for type: %s; request not present in async storage',
        this.type
      );

      return;
    }

    if (!this.exclusions) return;

    const { pathname } = parseurl(this.req);

    this.inputExclusions = this.exclusions
      .getInputExclusions('assess')
      .reduce((acc, e) => {
        if (!e.matchesUrl(pathname)) return acc;

        // `isNamed` is false for exclusion types such as BODY and QUERYSTRING which
        // apply to the entire set of params not only those by a specified name. Also
        // is true if the input name is set to wildcard "*" meaning it should apply to
        // all values. In each case there's no need to wrap if all rules are excluded.
        if (!this.skipEvent && e.appliesToAllAssessRules() && !e.isNamed) {
          logExclusionMessage(e, this.type);
          this.skipEvent = true;
        }

        acc.push(e);
        return acc;
      }, []);

    this.urlExclusions = this.exclusions
      .getUrlExclusions('assess')
      .reduce((acc, e) => {
        if (!e.matchesUrl(pathname)) return acc;

        if (!this.skipEvent && e.appliesToAllAssessRules()) {
          logExclusionMessage(e, this.type);
          this.skipEvent = true;
        }

        acc.push(e);
        return acc;
      }, []);
  }

  /**
   *
   * @param {object} param
   * @param {object} param.obj usually the incoming message
   * @param {string} param.prop obj[prop] is containing user-controlled data
   */
  handle({ obj, prop }) {
    this.typeKey = prop;

    if (this.skipEvent) return;

    this.traverseAndTrack({ obj, key: prop });
  }

  // eslint-disable-next-line complexity
  traverseAndTrack({ obj, key, path = [] }) {
    const value = obj[key];

    if (!value) return;

    if (this.sourceEventCount > SOURCE_EVENT_MAX) {
      logger.debug('max sources exceeded for %s', this.type);
      return;
    }

    if (Array.isArray(value)) {
      const limit = Math.min(value.length, this.samplingThreshold);

      trackedObjects.set(value, { ...this, path, key });

      for (let i = 0; i < limit; i += this.samplingInterval) {
        this.traverseAndTrack({
          obj: value,
          key: i,
          path: [...path, i]
        });
      }
    } else if (Buffer.isBuffer(value)) {
      this.sourceEventCount++;
      patcher.patch(value, 'toString', {
        name: 'Buffer.toString',
        alwaysRun: true,
        patchType: PATCH_TYPES.MISC,
        post: (wrapCtx) => {
          this.trackStringProp({
            obj: wrapCtx,
            key: 'result',
            path
          });
        }
      });
    } else if (typeof value === 'string' || value instanceof String) {
      this.sourceEventCount++;
      this.trackStringProp({ obj, key, path });
    } else if (typeof value === 'object') {
      trackedObjects.set(value, { ...this, path, key });

      for (const objKey of Object.keys(value)) {
        this.traverseAndTrack({
          obj: value,
          path: [...path, objKey],
          key: objKey
        });
      }
    }
  }

  trackStringProp({ obj, key, path, value = obj[key] }) {
    const trackData = tracker.track(value);
    if (!trackData) {
      return;
    }

    const name = path.join('.');
    const { props, str: result } = trackData;

    props.tagRanges = this.getTagRanges({ name, result });
    props.event = new SourceEvent({
      context: CallContext.create({
        obj: 'IncomingMessage {}',
        args: [`${this.typeKey}.${name}`],
        result,
        stackOpts: this.stackOpts
      }),
      code: `req.${this.typeKey}.${name}`,
      signature: this.signature,
      tagRanges: props.tagRanges,
      target: 'R',
      type: this.type,
      name
    });

    // NOTE: this used to happen on access in SourceMembrane.onString(), though we
    // should be able to determine access of source value during dataflow - when a
    // source is tracked into PROPAGATION or SINK event.
    storeSource({ type: this.type, name: key });

    obj[key] = result;
  }

  getTagRanges({ result, name }) {
    const stop = result.length - 1;
    const tagRanges = [new TagRange(0, stop, 'untrusted')];

    const typeLowerCase = this.type.toLowerCase();
    const nameLowerCase = name.toLocaleLowerCase();

    if (typeLowerCase === 'header' && nameLowerCase !== 'referer') {
      tagRanges.push(new TagRange(0, stop, 'header'));
    }

    if (typeLowerCase === 'cookie') {
      tagRanges.push(new TagRange(0, stop, 'cookie'));
    }

    if (!this.inputExclusions.length) {
      return tagRanges;
    }

    const rules = new Set();

    /*
      Maybe it's pointless because we set skipEvent to true if appliesToAllAssessRules
      and we won't get here
    */
    const exclusions = this.inputExclusions.filter(
      (e) => !e.appliesToAllAssessRules()
    );

    for (const exclusion of exclusions) {
      if (!exclusion.matches(name)) {
        continue;
      }

      for (const ruleId of exclusion.assessmentRulesList) {
        logger.debug(
          'excluding %s %s value from %s (%s)',
          this.type,
          name,
          ruleId,
          exclusion.name
        );

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

        rules.add(ruleId);
      }
    }

    return tagRanges;
  }
}

function storeSource({ type, name }) {
  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.
    storedSources = new Map();
    AsyncStorage.set(KEYS.SOURCES, storedSources);
  }

  // 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 }));
  }
}

function logExclusionMessage(exclusion, type) {
  const { name } = exclusion;
  logger.debug(
    'excluding %s inputs from all assess rules (%s)',
    type.toLowerCase(),
    name
  );
}

module.exports.SourceEventHandler = SourceEventHandler;
module.exports.trackedObjects = trackedObjects;
