javascript-obfuscator/src/node-transformers/converting-transformers/TemplateLiteralTransformer.ts

import { inject, injectable, } from 'inversify';
import { ServiceIdentifiers } from '../../container/ServiceIdentifiers';

import * as estraverse from 'estraverse';
import * as ESTree from 'estree';

import { TNodeWithScope } from '../../types/node/TNodeWithScope';
import { TStatement } from '../../types/node/TStatement';

import { IOptions } from '../../interfaces/options/IOptions';
import { IRandomGenerator } from '../../interfaces/utils/IRandomGenerator';
import { IVisitor } from '../../interfaces/node-transformers/IVisitor';

import { TransformationStage } from '../../enums/node-transformers/TransformationStage';

import { AbstractNodeTransformer } from '../AbstractNodeTransformer';
import { NodeFactory } from '../../node/NodeFactory';
import { NodeGuards } from '../../node/NodeGuards';
import { NodeUtils } from '../../node/NodeUtils';

/**
 * Transform ES2015 template literals to ES5
 * Thanks to Babel for algorithm
 */
@injectable()
export class TemplateLiteralTransformer extends AbstractNodeTransformer {
    /**
     * @param {IRandomGenerator} randomGenerator
     * @param {IOptions} options
     */
    constructor (
        @inject(ServiceIdentifiers.IRandomGenerator) randomGenerator: IRandomGenerator,
        @inject(ServiceIdentifiers.IOptions) options: IOptions
    ) {
        super(randomGenerator, options);
    }

    /**
     * @param {NodeGuards} node
     * @returns {boolean}
     */
    private static isLiteralNodeWithStringValue (node: ESTree.Node): boolean {
        return node && NodeGuards.isLiteralNode(node) && typeof node.value === 'string';
    }

    /**
     * @param {TransformationStage} transformationStage
     * @returns {IVisitor | null}
     */
    public getVisitor (transformationStage: TransformationStage): IVisitor | null {
        switch (transformationStage) {
            case TransformationStage.Converting:
                return {
                    enter: (node: ESTree.Node, parentNode: ESTree.Node | null) => {
                        if (parentNode && NodeGuards.isReturnStatementNode(node) && node.argument === null) {
                            return this.fixEsprimaReturnStatementTemplateLiteralNode(node);
                        }
                    },
                    leave: (node: ESTree.Node, parentNode: ESTree.Node | null) => {
                        if (parentNode && NodeGuards.isTemplateLiteralNode(node)) {
                            return this.transformNode(node, parentNode);
                        }
                    }
                };

            default:
                return null;
        }
    }

    /**
     * @param {TemplateLiteral} templateLiteralNode
     * @param {NodeGuards} parentNode
     * @returns {NodeGuards}
     */
    public transformNode (templateLiteralNode: ESTree.TemplateLiteral, parentNode: ESTree.Node): ESTree.Node {
        const templateLiteralExpressions: ESTree.Expression[] = templateLiteralNode.expressions;

        let nodes: ESTree.Expression[] = [];

        templateLiteralNode.quasis.forEach((templateElement: ESTree.TemplateElement) => {
            nodes.push(NodeFactory.literalNode(templateElement.value.cooked));

            const expression: ESTree.Expression | undefined = templateLiteralExpressions.shift();

            if (!expression) {
                return;
            }

            nodes.push(expression);
        });

        nodes = nodes.filter((node: ESTree.Literal | ESTree.Expression) => {
            return !(NodeGuards.isLiteralNode(node) && node.value === '');
        });

        // since `+` is left-to-right associative
        // ensure the first node is a string if first/second isn't
        if (
            !TemplateLiteralTransformer.isLiteralNodeWithStringValue(nodes[0]) &&
            !TemplateLiteralTransformer.isLiteralNodeWithStringValue(nodes[1])
        ) {
            nodes.unshift(NodeFactory.literalNode(''));
        }

        if (nodes.length > 1) {
            let root: ESTree.BinaryExpression = NodeFactory.binaryExpressionNode(
                '+',
                <ESTree.Literal>nodes.shift(),
                <ESTree.Expression>nodes.shift()
            );

            nodes.forEach((node: ESTree.Literal | ESTree.Expression) => {
                root = NodeFactory.binaryExpressionNode('+', root, <ESTree.Literal | ESTree.Expression>node);
            });

            return root;
        }

        return nodes[0];
    }

    /**
     * @param {ReturnStatement} returnStatementNode
     * @returns {Node | VisitorOption}
     */
    private fixEsprimaReturnStatementTemplateLiteralNode (returnStatementNode: ESTree.ReturnStatement): ESTree.Node | void {
        const scopeNode: TNodeWithScope = NodeUtils.getScopeOfNode(returnStatementNode);
        const scopeBody: TStatement[] = !NodeGuards.isSwitchCaseNode(scopeNode)
            ? scopeNode.body
            : scopeNode.consequent;
        const indexInScope: number = scopeBody.indexOf(returnStatementNode);

        // in incorrect AST-tree return statement node should be penultimate
        if (indexInScope !== scopeBody.length - 2) {
            return;
        }

        const nextSiblingStatementNode: TStatement | null = scopeBody[indexInScope + 1];

        if (!nextSiblingStatementNode || !NodeGuards.isExpressionStatementNode(nextSiblingStatementNode)) {
            return;
        }

        let isSiblingStatementHasTemplateLiteralNode: boolean = false;

        estraverse.traverse(nextSiblingStatementNode, {
            enter: (node: ESTree.Node, parentNode: ESTree.Node | null): void | estraverse.VisitorOption => {
                if (!NodeGuards.isTemplateLiteralNode(node)) {
                    return;
                }

                isSiblingStatementHasTemplateLiteralNode = true;

                return estraverse.VisitorOption.Break;
            }
        });

        if (!isSiblingStatementHasTemplateLiteralNode) {
            return;
        }

        returnStatementNode.argument = nextSiblingStatementNode.expression;
        scopeBody.pop();

        if (!NodeGuards.isSwitchCaseNode(scopeNode)) {
            scopeNode.body = [...scopeBody];
        } else {
            scopeNode.consequent = <ESTree.Statement[]>[...scopeBody];
        }

        return returnStatementNode;
    }
}