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
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
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
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
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
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
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
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)
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
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
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
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
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
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
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
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
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
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
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
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
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
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
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
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
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
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
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