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