/**
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 { get } = require('lodash');
const agentEmitter = require('../../agent-emitter');
const logger = require('../../core/logger')('contrast:hooks:restify');
const RestifyRequest = require('../../reporter/models/frameworks/restify-request');
const { Scopes, SCOPE_NAMES } = require('../../core/async-storage/scopes');
const moduleHook = require('../../hooks/require');
const patcher = require('../../hooks/patcher');
const { PATCH_TYPES } = require('../../constants');
const registerSessionHandler = require('./session');
const {
  constants: { EVENTS },
  utils
} = require('../../core/restify');

const {
  setFrameworkRequest,
  prototype: { decorateRequest }
} = require('../../hooks/frameworks/base');

class RestifyAssessSources {
  /**
   * @param {Boolean} [instrument=true] used for testing only
   */
  constructor(instrument = true) {
    if (instrument) {
      this.registerRestifyListeners();
      this.instrumentRestifyCookies();
    }
  }

  /**
   * Registers listeners for core restify events that will instrument
   * methods on restify framework
   */
  registerRestifyListeners() {
    agentEmitter.on(
      EVENTS.CREATE_SERVER,
      this.registerServerHandlers.bind(this)
    );
    agentEmitter.on(
      EVENTS.RESTIFY_EXPORT,
      this.registerQueryParamsHandler.bind(this)
    );
    agentEmitter.on(EVENTS.RESTIFY_EXPORT, this.registerBodyHandler.bind(this));
  }

  /**
   * Registers handlers for headers, url params and route coverage
   * @param {restify.Server} server
   */
  registerServerHandlers(server) {
    this.registerHeadersHandler(server);
    this.registerUrlParamsHandler(server);
  }

  /**
   * Register a `pre` handler to publish Assess event for header proxying.
   * @param {restify.Server} server
   */
  registerHeadersHandler(server) {
    utils.registerSafeHandler(server, 'pre', function(req, res, next) {
      setFrameworkRequest(req, RestifyRequest);
      agentEmitter.emit('assess.headers', req, 'headers');
    });
  }

  /**
   * The `bodyParser` plugin's factory function returns an Array of middleware
   * functions. This will patch the function to insert an additional 3rd handler
   * which will publish the assess event for body proxying after first 2 run.
   * @param {object} restify
   */
  registerBodyHandler(restify) {
    const {
      plugins: { bodyParser: bodyParserOrig }
    } = restify;

    if (bodyParserOrig) {
      patcher.patch(restify.plugins, 'bodyParser', {
        alwaysRun: true,
        name: 'restify.bodyParser',
        patchType: PATCH_TYPES.ASSESS_SOURCE,
        scopes: SCOPE_NAMES.ALLOW_ALL,
        post(data) {
          data.result.push(
            utils.makeSafeHandler((req) => {
              decorateRequest({ body: req.body });

              utils.emitMappedParams({
                options: get(data, 'args.0'),
                req
              });

              agentEmitter.emit('assess.body', req, 'body');
            })
          );
        }
      });
    }
  }

  /**
   * Registers a middleware after `server.use` that emits the `assess.params`
   * event to wrap in `SourceMembrane`
   * @param {restify.Server} server
   */
  registerUrlParamsHandler(server) {
    utils.registerSafeHandler(server, 'use', function(req, res, next) {
      decorateRequest({ normalizedUri: get(req, 'route.path') });
      utils.emitParams({ req });
    });
  }

  /**
   * Patches query parser middleware to wrap the callback argument. Once this
   * wrapped callback is invoked w/out error we know we can proxy `req.query`.
   * @param {object} restify
   */
  registerQueryParamsHandler(restify) {
    const {
      plugins: { queryParser: origQueryParser }
    } = restify;

    if (origQueryParser) {
      patcher.patch(restify.plugins, 'queryParser', {
        name: 'restify.plugins',
        patchType: PATCH_TYPES.FRAMEWORK,
        alwaysRun: true,
        post(data) {
          const [options] = data.args;
          patcher.patch(data, 'result', {
            name: 'restify.plugins.queryParser',
            patchType: PATCH_TYPES.ASSESS_SOURCE,
            alwaysRun: true,
            pre(data) {
              data.originalNext = data.args[2];
              data.args[2] = (...nextArgs) => {
                data.nextArgs = nextArgs;
              };
            },
            post(data) {
              if (data.nextArgs && !data.nextArgs.length) {
                try {
                  const [req] = data.args;
                  utils.emitMappedParams({ options, req });
                  agentEmitter.emit('assess.query', req, 'query');
                } catch (error) {
                  logger.error('error handling Assess query source: %o', error);
                }
              }

              Scopes.runInAllowAllScope(() => {
                data.originalNext(...data.nextArgs);
              });
            }
          });
        }
      });
    }
  }

  /**
   * Registers listener for require(`restify-cookies`) to patch middleware fn.
   */
  instrumentRestifyCookies() {
    moduleHook.resolve({ name: 'restify-cookies' }, (restifyCookies) => {
      this.registerCookieHandler(restifyCookies);
    });
  }

  /**
   * Patches restify-cookies to wrap the middleware's callback argument. Once
   * this wrapped callback is invoked w/out error we know we can proxy
   * `req.cookies`.
   * @param {import('restify-cookies')} restifyCookies
   */
  registerCookieHandler(restifyCookies) {
    patcher.patch(restifyCookies, 'parse', {
      alwaysRun: true,
      name: 'restify-cookies',
      patchType: PATCH_TYPES.ASSESS_SOURCE,
      pre(data) {
        data.originalNext = data.args[2];
        data.args[2] = (...nextArgs) => {
          data.nextArgs = nextArgs;
        };
      },
      post(data) {
        const [req, res] = data.args;

        registerSessionHandler(res);

        if (data.nextArgs && !data.nextArgs.length) {
          try {
            agentEmitter.emit('assess.cookies', req, 'cookies');
          } catch (error) {
            logger.error('error handling Assess cookie source: %o', error);
          }
        }

        Scopes.runInAllowAllScope(() => {
          data.originalNext(...data.nextArgs);
        });
      }
    });
  }
}

module.exports = RestifyAssessSources;
