/**
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 AsyncStorage = require('../async-storage');
const moduleHook = require('../../hooks/require');
const patcher = require('../../hooks/patcher');
const agentEmitter = require('../../agent-emitter');
const utils = require('./utils');
const {
  prototype: { decorateRequest },
  setFrameworkRequest
} = require('../../hooks/frameworks/base');
const KoaRequest = require('../../reporter/models/frameworks/koa-request');
const Stream = require('stream');
const tracker = require('../../tracker');
const { get } = require('lodash');
const { PATCH_TYPES } = require('../../constants');

const constants = {
  EVENTS: {
    KOA_EXPORT: 'koa-export', // used to patch methods on the koa framework
    CTX: 'koa-request-ctx', // used to patch the koa.ctx on _every_ request
    RES_END: 'koa-response-end', // used to handle ctx.res.end
    REQ_START: 'koa-request-start', // used to indicate a route is being handled
    CTX_REQUEST_BODY: 'koa-ctx-request-body', // used to handle when ctx.request.body is set
    CTX_REQUEST_HEADER: 'koa-ctx--request-header', // used to handle when ctx.request.header is set
    CTX_PARAMS: 'koa-ctx-params', // used to handle when ctx.params is set
    CTX_BODY: 'koa-ctx-body', // used when ctx.body is set, used for capture stack for xss assess rule
    MULTI_PART: 'koa-multer-multipart', // used when our multer wrapper middleware is called to emit the parsed req & res
    COOKIE: 'koa-cookie' // used when our cookie wrapper middleware is called to emit the parsed cookies
  }
};

class KoaCore {
  constructor() {
    this.instrumentKoa();
  }

  /**
   * Sets up handlers when `koa` and `koa-multer` and `koa-cookie` is required
   */
  instrumentKoa() {
    moduleHook.resolve({ name: 'koa' }, this.registerXportHandlers.bind(this));
    moduleHook.resolve(
      { name: 'koa-multer' },
      this.registerMultipartHandler.bind(this)
    );
    moduleHook.resolve(
      { name: 'koa-cookie' },
      this.registerCookieHandler.bind(this)
    );
    moduleHook.resolve(
      { name: 'http' },
      this.registerHttpResEndHandler.bind(this)
    );
  }

  /**
   * Sets up handlers when koa is required
   *
   * @param {Koa} Koa export prototype
   */
  registerXportHandlers(Koa) {
    this.registerHandleRequestHandler(Koa);
    this.registerMiddlewareHandler(Koa);
    this.registerCreateContextHandler(Koa);
  }

  /**
   * Wraps `koa-multer` to properly emit when a multipart form is parsed
   *
   * @param {Function} origMulter original multer middleware
   */
  registerMultipartHandler(origMulter) {
    return function multer(...args) {
      const instance = origMulter.apply(this, args);
      const origMake = instance._makeMiddleware;
      instance._makeMiddleware = function _makeMiddleware(...args) {
        const origMulterMiddleware = origMake.apply(this, args);

        return function multerMiddleware(req, res, origNext) {
          const next = function(...args) {
            // Attach the parsed body koa-multer places on `req` onto the the Koa
            // `request` object.
            // See: https://github.com/koa-modules/multer/pull/15/files
            req.request.body = req.body;
            agentEmitter.emit(constants.EVENTS.MULTI_PART, { req, res });
            return origNext.apply(this, args);
          };

          return origMulterMiddleware.apply(this, [req, res, next]);
        };
      };

      return instance;
    };
  }

  /**
   * Wraps `koa-cookie` to properly emit when cookies are parsed
   *
   * @param {Function} origCookie original cookie middleware
   */
  registerCookieHandler(origCookie) {
    function contrastCookie(...args) {
      const fn = origCookie.default.apply(this, args);

      return function contrastKoaCookie(ctx, next) {
        const contrastNext = function(...args) {
          agentEmitter.emit(constants.EVENTS.COOKIE, ctx);
          return next.apply(this, args);
        };

        return fn.apply(this, [ctx, contrastNext]);
      };
    }

    // this is exported on default in koa-cookie
    // https://github.com/varunpal/koa-cookie/blob/master/index.js#L30
    return {
      default: contrastCookie
    };
  }

  /**
   * patches `koa.prototype.handleRequest`
   * emits the `CTX` event for other handlers to instrument further
   * and stores the ctx in async storage
   * also patches `ctx.res.end` on `handleRequest` and emits `RES_END`
   *
   * @param {Koa} Koa export prototype
   */
  registerHandleRequestHandler(Koa) {
    patcher.patch(Koa.prototype, 'handleRequest', {
      name: 'Koa.Application',
      patchType: PATCH_TYPES.FRAMEWORK,
      alwaysRun: true,
      pre(data) {
        const [ctx] = data.args;
        agentEmitter.emit(constants.EVENTS.CTX, ctx);
        AsyncStorage.set(AsyncStorage.KEYS.KOA_CTX, ctx);
      }
    });
  }

  registerHttpResEndHandler(xport) {
    patcher.patch(xport.ServerResponse.prototype, 'end', {
      name: 'Koa res.end',
      patchType: PATCH_TYPES.SINK,
      alwaysRun: true,
      pre(data) {
        const ctx = AsyncStorage.get(AsyncStorage.KEYS.KOA_CTX);
        // this hook gets registered regardless of whether
        // koa is loaded - so need to check here
        // TODO conditionally load only if koa is loaded
        if (!ctx) {
          return;
        }
        // Koa supports body as a pipe, we want to ignore this
        if (!(ctx.body instanceof Stream)) {
          agentEmitter.emit(constants.EVENTS.RES_END, ctx);
        }
      }
    });
  }

  /**
   * Adds our middleware as first in a given request
   *
   * @param {Koa} Koa export prototype
   */
  registerMiddlewareHandler(Koa) {
    const self = this;
    function contrastStartMiddleware(ctx, next) {
      agentEmitter.emit(constants.EVENTS.REQ_START, ctx);
      self.registerRequestHandler(ctx);
      return next();
    }

    // mark these middleware as ours
    contrastStartMiddleware._isContrastStartMiddleware = true;

    patcher.patch(Koa.prototype, 'use', {
      name: 'Koa.Application',
      patchType: PATCH_TYPES.FRAMEWORK,
      alwaysRun: true,
      pre({ obj: app }) {
        // if not already inserted, insert the initial middleware.
        if (
          app.middleware &&
          (!app.middleware[0] || !app.middleware[0]._isContrastStartMiddleware)
        ) {
          app.middleware.splice(0, 0, contrastStartMiddleware);
        }
      }
    });
  }

  /**
   * Wraps ctx.request in a proxy to add set traps for `body`
   * and `header` to emit events for other instrumentation to handle
   *
   * @param {Koa.Context} ctx koa context instance
   */
  registerRequestHandler(ctx) {
    ctx.req.request = ctx.request;
    const { query } = ctx.request;
    ctx.request = new Proxy(ctx.request, {
      set(...args) {
        const [obj, prop] = args;
        const result = Reflect.set(...args);
        if (query) {
          decorateRequest({
            query
          });
        }
        switch (prop) {
          case 'body':
            decorateRequest({
              body: obj[prop]
            });
            agentEmitter.emit(constants.EVENTS.CTX_REQUEST_BODY, {
              request: obj,
              ctx
            });
            break;
          case 'header':
            agentEmitter.emit(constants.EVENTS.CTX_REQUEST_HEADER, {
              request: obj
            });
            break;
        }
        return result;
      }
    });
  }

  /**
   * patches `koa.prototype.createContext`, creates a proxy
   * for the `ctx` to add set traps for `params` and `body`
   * to emit events for further instrumentation
   *
   * @param {Koa} Koa export prototype
   */
  registerCreateContextHandler(Koa) {
    patcher.patch(Koa.prototype, 'createContext', {
      name: 'Koa.Application',
      patchType: PATCH_TYPES.FRAMEWORK,
      alwaysRun: true,
      post(wrapCtx) {
        // Set the framework-specific request model in our first hook point
        setFrameworkRequest(wrapCtx.result, KoaRequest);
        // ctx.params get attached to context with koa-router, this is the
        // easiest way to get at this data
        // for ctx.body we want a relevant stack trace so this is where we do it
        wrapCtx.result = new Proxy(wrapCtx.result, {
          set(...args) {
            const [ctx, prop, value] = args;
            const result = Reflect.set(...args);
            let props; // have to define here because eslint complains about `no-case-declarations`
            switch (prop) {
              case 'params':
                decorateRequest({
                  parameters: ctx.params,
                  normalizedUri: ctx._matchedRoute
                });

                agentEmitter.emit(constants.EVENTS.CTX_PARAMS, ctx);
                break;
              case 'body':
                props = tracker.getData(value);
                agentEmitter.emit(constants.EVENTS.CTX_BODY, {
                  ctx: wrapCtx.result,
                  tracked: get(props, 'tracked', false)
                });
            }
            return result;
          }
        });
      }
    });
  }
}

module.exports = KoaCore;
module.exports.constants = constants;
module.exports.utils = utils;
