Ejemplo n.º 1
0
def opt_convert_type_to_char(ast):
    for node in ast.iter():
        if node.tag in ["ConvertExpressionAst"]:
            type_name = node.find("TypeConstraintAst")
            if type_name is not None:
                type_name = type_name.attrib["TypeName"].lower()

            if type_name == "char":
                cst_int_node = node.find("ConstantExpressionAst")

                if cst_int_node is not None and cst_int_node.attrib[
                        "StaticType"] == "int":
                    type_value = int(cst_int_node.text)

                    new_element = Element(
                        "StringConstantExpressionAst", {
                            "StringConstantType": "SingleQuoted",
                            "StaticType": "string",
                        })
                    new_element.text = chr(type_value)

                    log_debug("Replace integer %d convertion to char '%s'" %
                              (type_value, new_element.text))

                    replace_node(ast, node, new_element)

                    return True

    return False
Ejemplo n.º 2
0
def opt_constant_string_type(ast):
    for node in ast.iter():
        if node.tag in ["InvokeMemberExpressionAst", "MemberExpressionAst"]:
            for cst_string_node in node.findall("StringConstantExpressionAst"):
                member = cst_string_node.text.lower()
                if member in BAREWORDS:
                    if cst_string_node.attrib["StringConstantType"] != "BareWord":
                        cst_string_node.text = BAREWORDS[member]

                        log_debug("Fix member string type for '%s'" % cst_string_node.text)

                        cst_string_node.attrib["StringConstantType"] = "BareWord"

                        return True

        if node.tag in ["CommandElements"]:
            for subnode in node:
                if subnode.tag == "StringConstantExpressionAst" and subnode.attrib["StringConstantType"] != "BareWord":
                    subnode.attrib["StringConstantType"] = "BareWord"

                    log_debug("Fix command string type for '%s'" % subnode.text)

                    return True
                break
    return False
Ejemplo n.º 3
0
def opt_binary_expression_format(ast):
    for node in ast.iter():
        if node.tag in ["BinaryExpressionAst"
                        ] and node.attrib["Operator"] == "Format":
            format_str = node.find("StringConstantExpressionAst")
            if format_str is not None:
                format_str = format_str.text

            argument_values = get_array_literal_values(
                node.find("ArrayLiteralAst"))
            if argument_values is None:
                continue

            try:
                formatted = format_str.format(*argument_values)
            except IndexError:
                continue

            new_element = Element("StringConstantExpressionAst", {
                "StringConstantType": "SingleQuoted",
                "StaticType": "string",
            })
            new_element.text = formatted

            log_debug("Apply format operation to '%s'" % formatted)

            replace_node(ast, node, new_element)

            return True
    return False
Ejemplo n.º 4
0
def opt_invoke_expression(ast):
    ret = False
    p = pathlib.Path("tmp.ps1")

    for node in ast.iter():
        if node.tag == "CommandElements":
            subnodes = list(node)
            if len(subnodes) == 2:
                if subnodes[0].tag == "StringConstantExpressionAst" and subnodes[0].attrib["StringConstantType"] == "BareWord" and subnodes[0].text == "Invoke-Expression":
                    if subnodes[1].tag == "StringConstantExpressionAst" and subnodes[1].attrib["StringConstantType"] != "BareWord":

                        script_content = subnodes[1].text

                        with open(p, "w") as tmp:
                            tmp.write(script_content)

                        if create_ast_file(p):
                            if sub_ast := read_ast_file(p.with_suffix(".xml")):

                                log_debug("Replace Invoke-Expression by expression AST")

                                replace_node(ast, subnodes[0], sub_ast.getroot(), until="CommandAst")

                                ret = True
                                break
Ejemplo n.º 5
0
def opt_binary_expression_replace(ast):
    for node in ast.iter():
        if node.tag in ["BinaryExpressionAst"
                        ] and node.attrib["Operator"] == "Ireplace":
            target = node.find("StringConstantExpressionAst")
            if target is not None:
                target = target.text

            argument_values = get_array_literal_values(
                node.find("ArrayLiteralAst"))

            if argument_values is None or len(argument_values) != 2:
                return False

            formatted = target.replace(argument_values[0], argument_values[1])

            log_debug("Apply replace operator to '%s'" % formatted)

            new_element = Element("StringConstantExpressionAst", {
                "StringConstantType": "SingleQuoted",
                "StaticType": "string",
            })
            new_element.text = formatted

            replace_node(ast, node, new_element)

            return True
    return False
Ejemplo n.º 6
0
def opt_invoke_split_string(ast):
    for node in ast.iter():
        if node.tag == "InvokeMemberExpressionAst":
            subnodes = list(node)

            if len(subnodes) < 3:
                continue

            if subnodes[2].tag == 'StringConstantExpressionAst' and \
                    subnodes[2].attrib["StringConstantType"] == "BareWord" and \
                    subnodes[2].text.lower() == "split":
                if subnodes[1].tag == 'StringConstantExpressionAst' and \
                        subnodes[1].attrib["StringConstantType"] != "BareWord":
                    argument = subnodes[0]
                    if argument is not None:
                        argument = argument.find("StringConstantExpressionAst")
                        if argument is not None:

                            splitted = subnodes[1].text.split(argument.text)

                            new_array_ast = create_array_literal_values(splitted)

                            log_debug("Apply split operation to %s" % splitted)

                            replace_node(ast, node, new_array_ast)
                            return True
    return False
Ejemplo n.º 7
0
def try_reverse_variable_if_not_used(ast, variable, before_node):
    parent_map = dict((c, p) for p in ast.iter() for c in p)

    for node in ast.iter():
        if node.tag == "VariableExpressionAst" and node.attrib["VariablePath"].lower() == variable.lower():
            parent = parent_map[node]
            if parent is not None and parent_map[node].tag == "AssignmentStatementAst":
                operands = parent.find("CommandExpressionAst")
                if operands.tag == "CommandExpressionAst":
                    operands = operands.find("ArrayLiteralAst")
                if operands is not None:
                    operands = operands.find("Elements")

                    new_element = Element("Elements")
                    for element in operands:
                        new_element.insert(0, element)

                    replace_node(ast, operands, new_element)

                    log_debug(f"Apply reverse method to variable ${variable}")

                    return True
            else:
                return False

    return False
Ejemplo n.º 8
0
def opt_convert_type_to_array(ast):
    for node in ast.iter():
        if node.tag in ["ConvertExpressionAst"]:
            type_name = node.find("TypeConstraintAst")
            if type_name is not None:
                type_name = type_name.attrib["TypeName"].lower()

            if type_name == "array":
                cst_string_node = node.find("StringConstantExpressionAst")
                if cst_string_node is not None:

                    log_debug("Replace array of one string to string '%s'" %
                              cst_string_node.text)

                    replace_node(ast, node, cst_string_node)

            elif type_name == "char[]":
                cst_string_node = node.find("StringConstantExpressionAst")
                if cst_string_node is not None:
                    arrayed = [c for c in cst_string_node.text]

                    new_array_ast = create_array_literal_values(arrayed)

                    log_debug("Replace (cast) string to array: '%s'" % arrayed)

                    replace_node(ast, node, new_array_ast)
Ejemplo n.º 9
0
def opt_convert_type_to_type(ast):
    for node in ast.iter():
        if node.tag in ["ConvertExpressionAst"]:
            type_name = node.find("TypeConstraintAst")
            if type_name is not None:
                type_name = type_name.attrib["TypeName"].lower()

            if type_name in ["type"]:
                cst_string_node = node.find("StringConstantExpressionAst")
                if cst_string_node is not None:
                    type_value = cst_string_node.text

                    new_element = Element("StringConstantExpressionAst", {
                        "StringConstantType": "BareWord",
                        "StaticType": "string",
                    })

                    new_element.text = "[" + type_value + "]"

                    log_debug("Replace type string '%s' by type '%s'" %
                              (type_value, new_element.text))

                    replace_node(ast, node, new_element)

                    return True
Ejemplo n.º 10
0
def opt_value_of_const_array(ast):
    for node in ast.iter():
        if node.tag == "IndexExpressionAst":
            subnodes = list(node)

            if subnodes[0].tag == "StringConstantExpressionAst":
                target = subnodes[0].text
            elif subnodes[0].tag == "ArrayLiteralAst":
                target = get_array_literal_values(subnodes[0])
            else:
                continue

            if subnodes[1].tag == "ConstantExpressionAst":
                indexes = [int(subnodes[1].text)]
            elif subnodes[1].tag == "ArrayLiteralAst":
                values = get_array_literal_values(subnodes[1])
                indexes = values
            else:
                continue

            if target is not None and indexes is not None:
                if len(indexes) > 0:

                    new_array_ast = create_array_literal_values([target[index] for index in indexes])

                    log_debug(f"Apply index {indexes} operation to constant {target.__class__.__name__} {target}")

                    replace_node(ast, node, new_array_ast)

                    return True

    return False
Ejemplo n.º 11
0
def opt_binary_expression_plus(ast):
    for node in ast.iter():
        if node.tag == 'BinaryExpressionAst':
            operator = node.attrib['Operator']
            if operator == "Plus":
                subnodes = list(node)

                if subnodes[0].tag == "StringConstantExpressionAst":
                    left = subnodes[0].text
                elif subnodes[0].tag == "ArrayLiteralAst":
                    left = get_array_literal_values(subnodes[0])
                else:
                    continue

                if subnodes[1].tag == "StringConstantExpressionAst":
                    right = subnodes[1].text
                elif subnodes[1].tag == "ArrayLiteralAst":
                    right = get_array_literal_values(subnodes[1])
                else:
                    continue

                if left is not None and right is not None:
                    if isinstance(left, str) and isinstance(right, str):

                        new_element = Element('StringConstantExpressionAst')
                        new_element.set('StringConstantType', 'DoubleQuoted')
                        new_element.text = left + right

                        log_debug(
                            "Merging constant strings:  '%s', '%s' to '%s'" %
                            (subnodes[0].text, subnodes[1].text,
                             new_element.text))

                        replace_node(ast, node, new_element)

                        return True

                    else:
                        items = []
                        if isinstance(left, str) and isinstance(right, list):
                            right.insert(0, left)
                            items = right
                        elif isinstance(left, list) and isinstance(right, str):
                            left.append(right)
                            items = left
                        elif isinstance(left, list) and isinstance(
                                right, list):
                            left.extend(right)
                            items = left

                        new_array_ast = create_array_literal_values(items)

                        replace_node(ast, node, new_array_ast)

                        return True
    return False
Ejemplo n.º 12
0
def opt_special_variable_case(ast):
    for node in ast.iter():
        if node.tag == "VariableExpressionAst":
            if node.attrib["VariablePath"].lower() in SPECIAL_VARS_NAMES:
                if node.attrib["VariablePath"] != SPECIAL_VARS_NAMES[node.attrib["VariablePath"].lower()]:
                    node.attrib["VariablePath"] = SPECIAL_VARS_NAMES[node.attrib["VariablePath"].lower()]

                    log_debug(f'Fix variable name case for ${node.attrib["VariablePath"]}')

                    return True
    return False
Ejemplo n.º 13
0
def opt_command_element_as_bareword(ast):
    for node in ast.iter():
        if node.tag == "CommandElements":
            for subnode in node:
                if subnode.tag == "StringConstantExpressionAst" and subnode.attrib["StringConstantType"] != "BareWord":
                    if subnode.text in BAREWORDS:
                        subnode.attrib["StringConstantType"] = "BareWord"

                        log_debug(f"Fix string type for command {subnode.text}")

                        return True
    return False
Ejemplo n.º 14
0
def create_ast_file(ps1_file):
    log_info(f"Creating AST for: {ps1_file}")

    cmd = ["PowerShell", "-ExecutionPolicy", "Unrestricted", "-File",
           os.path.abspath(os.path.join("tools", "Get-AST.ps1")),
           "-ps1", os.path.abspath(ps1_file)]

    result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)

    for line in result.stdout.splitlines():
        log_debug(line)

    return result.returncode == 0
Ejemplo n.º 15
0
def opt_remove_empty_nodes(ast):
    for node in ast.iter():
        if node.tag in ["Attributes", "Redirections", "CatchTypes"]:
            subnodes = list(node)
            if len(subnodes) == 0:

                log_debug(f"Remove empty node {node.tag}")

                delete_node(ast, node)

                return True

    return False
Ejemplo n.º 16
0
def opt_simplify_pipeline_single_command(ast):
    for node in ast.iter():
        if node.tag == "PipelineAst":
            subnodes = list(node)
            if len(subnodes) == 1 and subnodes[0].tag in ["PipelineElements"]:
                subnodes = list(subnodes[0])
            if len(subnodes) == 1:

                log_debug("Replace pipeline with single elements by %s" % subnodes[0].tag)

                replace_node(ast, node, subnodes[0])

                return True
    return False
Ejemplo n.º 17
0
def opt_replace_constant_variable_by_value(ast):
    cst_assigned = dict()

    used_vars = get_used_vars(ast)

    for node in ast.iter():
        if node.tag in ["AssignmentStatementAst"]:
            subnodes = list(node)
            if subnodes[0].tag == "VariableExpressionAst":
                variable = subnodes[0]
                if subnodes[1].tag == "CommandExpressionAst":
                    subnodes = list(subnodes[1])
                    if len(subnodes) == 1:
                        if subnodes[0].tag == "StringConstantExpressionAst":
                            cst_assigned[variable.attrib["VariablePath"].lower()] = subnodes[0].text
                        elif subnodes[0].tag == "ArrayLiteralAst":
                            cst_assigned[variable.attrib["VariablePath"].lower()] = get_array_literal_values(
                                subnodes[0])
                    else:
                        if variable.attrib["VariablePath"].lower() in cst_assigned:
                            del cst_assigned[variable.attrib["VariablePath"].lower()]

        if node.tag in ["UnaryExpressionAst", "BinaryExpressionAst", "Arguments", "InvokeMemberExpressionAst"]:
            subnodes = list(node)
            for subnode in subnodes:
                if subnode.tag == "VariableExpressionAst":
                    var_name = subnode.attrib["VariablePath"].lower()
                    if var_name in cst_assigned and used_vars.setdefault(var_name, 0) == 1:

                        value = cst_assigned[var_name]

                        if isinstance(value, str):
                            new_element = create_constant_string(value,
                                                                 "BareWord" if node.tag == "InvokeMemberExpressionAst"
                                                                 else "DoubleQuoted")

                            log_debug("Replace constant variable %s (string) in expression" % (
                                subnode.attrib["VariablePath"]))

                            replace_node(ast, subnode, new_element)
                            return True

                        elif isinstance(value, list):
                            new_element = create_array_literal_values(value)
                            log_debug(
                                "Replace constant variable %s (array) in expression" % (subnode.attrib["VariablePath"]))
                            replace_node(ast, subnode, new_element)
                            return True

    return False
Ejemplo n.º 18
0
    def optimize(self, ast):

        count_in = sum(1 for _ in ast.getroot().iter())
        log_debug(f"{count_in} nodes loaded")

        while optimize_pass(ast):
            self.stats.modifications += 1

        log_info(f"{self.stats.modifications} modifications applied")

        count_out = sum(1 for _ in ast.getroot().iter())
        ratio = "{:02.2f}".format(count_out / count_in * 100.00)
        log_debug(f"{count_out} nodes in output ({ratio}%)")

        return ast
Ejemplo n.º 19
0
def opt_simplify_single_array(ast):
    for node in ast.iter():
        if node.tag == "ArrayLiteralAst":
            subnodes = list(node)
            if len(subnodes) == 1 and subnodes[0].tag in ["Elements"]:
                subnodes = list(subnodes[0])
            if len(subnodes) == 1 and subnodes[0].tag not in ["CommandAst", "UnaryExpressionAst",
                                                              "BinaryExpressionAst"]:

                log_debug("Replace array with single element by %s" % subnodes[0].tag)

                replace_node(ast, node, subnodes[0])

                return True
    return False
Ejemplo n.º 20
0
def opt_prefixed_variable_case(ast):
    for node in ast.iter():
        if node.tag == "StringConstantExpressionAst" and node.attrib["StringConstantType"] == "BareWord":
            names = node.text.split(":")
            if len(names) > 1 and names[0].lower() in ["variable", "env"]:
                old_name = node.text
                names[0] = names[0].lower()
                new_name = ":".join(names)

                if old_name != new_name:
                    node.text = new_name

                    log_debug("Fix string case from '%s' to '%s'" % (old_name, node.text))

                    return True
    return False
Ejemplo n.º 21
0
def opt_alias(ast):
    for node in ast.iter():
        if node.tag in ["StringConstantExpressionAst"] and node.attrib["StringConstantType"] == "BareWord":
            old_value = node.text

            new_value = ALIAS[old_value.lower()] if old_value.lower() in ALIAS else old_value

            if old_value != new_value:

                log_debug("Replace alias of %s by %s" % (old_value, new_value))

                node.text = new_value

                return True

    return False
Ejemplo n.º 22
0
def opt_simplify_paren_single_expression(ast):
    for node in ast.iter():
        if node.tag == "ParenExpressionAst":
            subnodes = list(node)
            if len(subnodes) == 1 and subnodes[0].tag in ["PipelineAst"]:
                subnodes = list(subnodes[0])
            if len(subnodes) == 1 and subnodes[0].tag in ["PipelineElements"]:
                subnodes = list(subnodes[0])
            if len(subnodes) == 1 and subnodes[0].tag in ["CommandExpressionAst"]:
                subnodes = list(subnodes[0])
            if len(subnodes) == 1 and subnodes[0].tag not in ["CommandAst", "UnaryExpressionAst",
                                                              "BinaryExpressionAst"]:

                log_debug("Replace paren with single expression by %s" % subnodes[0].tag)

                replace_node(ast, node, subnodes[0])

                return True
    return False
Ejemplo n.º 23
0
def opt_type_constraint_case(ast):
    for node in ast.iter():
        if node.tag in ["TypeConstraintAst", "TypeExpressionAst"]:
            typename = node.attrib["TypeName"]

            new_value = typename
            new_value = ".".join(
                [BAREWORDS[t.lower()] if t.lower() in BAREWORDS else t for t in new_value.split(".")])
            new_value = "-".join(
                [BAREWORDS[t.lower()] if t.lower() in BAREWORDS else t for t in new_value.split("-")])

            if typename != new_value:
                node.attrib["TypeName"] = new_value

                log_debug("Fix typename case from '%s' to '%s'" % (typename, new_value))

                return True

    return False
Ejemplo n.º 24
0
def opt_binary_expression_join(ast):
    for node in ast.iter():
        if node.tag in ["BinaryExpressionAst"
                        ] and node.attrib["Operator"] == "Join":
            subnodes = list(node)

            joiner = node.find("StringConstantExpressionAst")
            if joiner is not None:
                joiner = joiner.text
                if joiner is None:
                    joiner = ""
            else:
                log_err(
                    f"BinaryExpression Join with {subnodes[0].tag} joiner is unsupported"
                )
                continue

            values = node.find("ArrayLiteralAst")
            if values is not None:
                values = get_array_literal_values(values)

            if joiner is None or values is None:
                continue

            try:
                joined = joiner.join(values)
            except Exception:
                continue

            new_element = Element("StringConstantExpressionAst", {
                "StringConstantType": "SingleQuoted",
                "StaticType": "string",
            })
            new_element.text = joined

            log_debug("Apply join operation to '%s'" % joined)

            replace_node(ast, node, new_element)

            return True
    return False
Ejemplo n.º 25
0
def opt_invoke_replace_string(ast):
    for node in ast.iter():
        if node.tag == "InvokeMemberExpressionAst":
            subnodes = list(node)

            if len(subnodes) < 3:
                continue

            if subnodes[2].tag == 'StringConstantExpressionAst' and \
                    subnodes[2].attrib["StringConstantType"] == "BareWord" and \
                    subnodes[2].text.lower() == "replace":
                if subnodes[1].tag == 'StringConstantExpressionAst' and \
                        subnodes[1].attrib["StringConstantType"] != "BareWord":
                    arguments = subnodes[0]
                    if arguments is not None:
                        argument_values = []

                        for element in list(arguments):
                            if element.tag == "StringConstantExpressionAst":
                                argument_values.append(element.text)

                        if len(argument_values) != 2:
                            continue

                        formatted = subnodes[1].text.replace(argument_values[0], argument_values[1])

                        log_debug("Apply replace method on '%s'" % formatted)

                        new_element = Element("StringConstantExpressionAst",
                                              {
                                                  "StringConstantType": "SingleQuoted",
                                                  "StaticType"        : "string",
                                              })
                        new_element.text = formatted

                        replace_node(ast, node, new_element)

                        return True
    return False
Ejemplo n.º 26
0
def opt_bareword_case(ast):
    for node in ast.iter():
        if node.tag in ["StringConstantExpressionAst"] and node.attrib["StringConstantType"] == "BareWord":
            old_value = node.text

            new_value = node.text
            is_type = new_value[0] == "[" and new_value[-1] == "]"
            if is_type:
                new_value = new_value[1:-1]

            new_value = ".".join([BAREWORDS[t.lower()] if t.lower() in BAREWORDS else t for t in new_value.split(".")])
            new_value = "-".join([BAREWORDS[t.lower()] if t.lower() in BAREWORDS else t for t in new_value.split("-")])

            if is_type:
                new_value = "[" + new_value + "]"

            if old_value != new_value:
                node.text = new_value

                log_debug("Fix bareword case from '%s' to '%s'" % (old_value, node.text))

                return True

    return False
Ejemplo n.º 27
0
def opt_convert_type_to_string(ast):
    for node in ast.iter():
        if node.tag in ["ConvertExpressionAst"]:
            type_name = node.find("TypeConstraintAst")
            if type_name is not None:
                type_name = type_name.attrib["TypeName"].lower()

            if type_name in ["string"]:
                cst_string_node = node.find("VariableExpressionAst")
                if cst_string_node is not None:
                    var_value = cst_string_node.attrib["VariablePath"]

                    if var_value.lower(
                    ) in SPECIAL_VARS_VALUES and SPECIAL_VARS_VALUES[
                            var_value.lower()] is not None:
                        log_debug("Use special variable value '%s' for $%s" %
                                  (SPECIAL_VARS_VALUES[var_value.lower()],
                                   var_value))
                        var_value = SPECIAL_VARS_VALUES[var_value.lower()]

                    new_element = Element(
                        "StringConstantExpressionAst", {
                            "StringConstantType": "DoubleQuoted",
                            "StaticType": "string",
                        })

                    new_element.text = var_value

                    log_debug("Replace type of variable $%s to string" %
                              (var_value))

                    replace_node(ast, node, new_element)

                    return True

                cst_string_node = node.find("StringConstantExpressionAst")
                if cst_string_node is not None:

                    log_debug("Remove unused cast to string for '%s'" %
                              (cst_string_node.text))

                    replace_node(ast, node, cst_string_node)

                    return True