/**
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 { default: generate } = require('@babel/generator');
const { parse } = require('@babel/parser');
const { default: traverse } = require('@babel/traverse');

const RewriteDeadzones = require('../../assess/deadzones/rewrite');
const { util: sourceMapUtility } = require('../../util/source-map');
const Scopes = require('../async-storage/scopes');
const logger = require('../logger')('contrast:rewrite');
const AssignmentExpression = require('./assignment-expression');
const BinaryExpression = require('./binary-expression');
const CallExpression = require('./call-expression');
const CatchClause = require('./catch-clause');
const functionWrap = require('./function-wrap');
// const ImportDeclaration = require('./import-declaration');
const isContrastMethod = require('./is-contrast-method');
const logRewrite = require('./log');
const MemberExpression = require('./member-expression');
const ObjectProperty = require('./object-property');
const prependGlobals = require('./prepend-globals');
const RewriteLog = require('./rewrite-log');
const SwitchStatement = require('./switch-statement');
const TemplateLiteral = require('./template-literal');

/**
 * @typedef {Object} State
 * @property {import('../../agent')} agent
 * @property {RewriteLog} rewriteLog
 * @property {import('./callees').Spec} specs
 * @property {{ [name: string]: import('@babel/types').MemberExpression }} callees
 * @property {boolean} inject
 * @property {boolean} wrap
 */

const DEADZONE_NAME = 'agent: rewrite';

/** @type {Rewriter} memoized rewriter return value. */
let rewriter;

class Rewriter {
  constructor(agent) {
    this.agent = agent;
    const { specs, callees } = require('./callees')(agent);
    this.specs = specs;
    this.callees = callees;
    this.deadzones = new RewriteDeadzones(agent);
  }

  /**
   * @param {string} code
   * @param {string} filename
   * @returns {import('@babel/types').File} AST
   */
  parse(code, filename, sourceType) {
    return parse(code, {
      // TODO: check if necessary
      plugins: [
        'classPrivateMethods',
        'classPrivateProperties',
        'classProperties'
      ],
      ranges: true,
      sourceType,
      sourceFilename: filename,
      tokens: true
    });
  }

  /**
   *
   * @param {import('@babel/types').File} ast
   * @param {string} filename
   * @param {State} initialState
   */
  traverse(ast, filename, initialState) {
    traverse(
      ast,
      {
        /** Runs before other visitor functions. */
        enter(path, state) {
          for (let i = 0; i < state.agent.staticVisitors.length; i++) {
            state.rewriteLog.logRewrite('staticVisitors');
            state.agent.staticVisitors[i](path.node, filename, state);
          }
        },
        Program(path, state) {
          if (state.wrap) functionWrap(path, state);
          if (state.inject) prependGlobals(path, state);
        },
        AssignmentExpression,
        BinaryExpression,
        CallExpression,
        CatchClause,
        // ImportDeclaration, disabled since we're not yet utilizing this.
        MemberExpression,
        ObjectProperty,
        SwitchStatement,
        TaggedTemplateExpression: TemplateLiteral.skip,
        TemplateLiteral
      },
      null,
      initialState
    );
    traverse.cache.clear();
  }

  /**
   * Transform original javascript into contrast rewritten form.
   * Should be run in no instrumentation scope
   *
   * @param {string} code
   * @param {string} filename
   * @param {State} initialState
   * @param {string} sourceType 'script' for cjs, 'module' for esm
   * @returns {import('@babel/generator').GeneratorResult}
   */
  rewrite(code, filename, initialState, sourceType = 'script') {
    // I don't think this changes but use the native module wrap to
    // code = require('module').wrap(code);
    // XXX this was causing an error by wrapping es6 features in a function
    if (!this.agent.config.agent.node.enable_rewrite) {
      return { code };
    }

    try {
      // Comment out shebang if present, won't affect any code locations.
      code = code.charAt(0) === '#' ? `//${code}` : code;

      const ast = this.parse(code, filename, sourceType);
      this.traverse(ast, filename, initialState);
      return generate(
        ast,
        {
          jsonCompatibleStrings: true,
          // source maps are generated regardless of whether they're
          // configured to be used in reporting. This has no slowdown
          // overhead except during startup, which already sucks.
          // use rewrite cache!
          sourceMaps: true,
          sourceFileName: filename
        },
        code
      );
    } catch (err) {
      initialState.rewriteLog.aborted = true;

      logger.warn('could not rewrite %s', filename);
      logger.warn(err);
    }
  }

  /**
   * Takes code string, creates an AST, traverses the AST to preform transformations we need.
   * @param {string} content
   * @param {string} filename
   * @param {Object} opts
   * @param {boolean} opts.inject whether or not to inject globals
   * @param {string} opts.sourceType 'script' for cjs, 'module' for esm
   * @param {boolean} opts.wrap whether or not to function wrap
   * @returns {import('@babel/generator').GeneratorResult}
   */
  rewriteCode(
    content,
    filename,
    { inject = true, sourceType = 'script', wrap = false } = {}
  ) {
    const code = String(content);

    /** @type {import('@babel/generator').GeneratorResult} */
    const result = { code };

    // Determine if we even need to rewrite: if there are no callees, no.
    if (this.callees && Object.keys(this.callees).length === 0) {
      logger.debug('no custom callees, skipping rewrite for %s', filename);
      return result;
    }

    const rewriteLog = new RewriteLog(this.specs, this.callees, code);

    // ...if there's no tokens to rewrite or static visitors to run, no.
    if (!rewriteLog.foundTokens()) {
      logger.debug('no rewritable tokens, skipping rewrite for %s', filename);
      return result;
    }

    /** @type {State} */
    const initialState = {
      agent: this.agent,
      rewriteLog,
      specs: this.specs,
      callees: this.callees,
      inject,
      wrap: wrap && sourceType === 'script',
      deps: []
    };

    const rewriteResult = Scopes.runInNoInstrumentationScope(
      () => this.rewrite(code, filename, initialState, sourceType),
      DEADZONE_NAME
    );

    if (rewriteLog.aborted) {
      logger.debug('rewrite aborted, returning original code for %s', filename);
      return result;
    }

    logger.trace('successfully rewrote code for %s', filename);
    rewriteResult.deps = initialState.deps;

    return rewriteResult;
  }

  /**
   * @param {string} content file content
   * @param {string} filename source filename
   * @param {Object} opts
   * @param {boolean} opts.inject whether or not to inject globals
   * @param {string} opts.sourceType 'script' for cjs, 'module' for esm
   * @param {boolean} opts.wrap whether or not to function wrap
   * @returns {import('@babel/generator').GeneratorResult}
   */
  rewriteFile(
    content,
    filename,
    { inject = true, sourceType = 'script', wrap = true } = {}
  ) {
    if (this.deadzones.isDeadzoned(filename)) {
      logger.debug('skipping rewrite for deadzoned file %s', filename);
      return content;
    }

    const result = this.rewriteCode(content, filename, {
      inject,
      sourceType,
      wrap
    });
    if (this.agent.config.agent.node.enable_rewrite_log) {
      logRewrite(result.code, filename);
    }

    // if sourceMapUtility is not configured, this will throw a null
    // exception. As such, this will almost always fail in tests.
    try {
      result.map = sourceMapUtility.chain(filename, result.map);
    } catch (err) {
      logger.debug(`Unable to chain source map for ${filename} %o`, err);
    }
    return result;
  }

  /**
   * Unwrites the global.contrastMethods references when a user calls
   * `Function.prototype.toString`.
   * @param {string} code
   * @returns {string} sanitized code
   */
  unwrite(code) {
    const ast = this.parse(code);

    const self = this;
    traverse(ast, {
      // All rewritten code to remove has been converted to a Call Expression...
      CallExpression(path) {
        if (!isContrastMethod(path.node.callee)) return;

        const spec = self.specs.find(
          (spec) => spec.name === path.node.callee.property.name
        );
        if (!spec) return;

        switch (spec.type) {
          case 'BinaryExpression': {
            return BinaryExpression.unwrite(path, spec);
          }
          case 'CallExpression': {
            return CallExpression.unwrite(path);
          }
          case 'MemberExpression': {
            return MemberExpression.unwrite(path);
          }
          case 'SwitchStatement': {
            return SwitchStatement.unwrite(path);
          }
        }
      },
      // ...or a Tagged Template Expression
      TaggedTemplateExpression: TemplateLiteral.unwrite
    });

    return generate(ast, { jsonCompatibleStrings: true }).code;
  }
}

/**
 * Returns a memoized result, since a rewriter is included in multiple locations.
 * @returns {Rewriter}
 */
module.exports = function createBabelRewriter(agent) {
  if (!rewriter) rewriter = new Rewriter(agent);
  return rewriter;
};
