def _check_for_iocs(loop, context, indent):
    """Generate Python JIT code for checking the variables modified in a
    loop to see if they were set to interesting IOCs.

    @param loop (VBA_Object object) The loop for which to generate
    Python JIT code.

    @param context (Context object) The current program state.

    @param indent (int) The number of spaces to indent the generated
    Python code.

    @return (str) Python JIT code checking variables modified in the
    loop for potential IOCs.

    """
    context = context # pylint
    
    indent_str = " " * indent
    lhs_visitor = lhs_var_visitor()
    loop.accept(lhs_visitor)
    lhs_var_names = lhs_visitor.variables
    ioc_str = indent_str + "# Check for IOCs in intermediate variables.\n"
    for var in lhs_var_names:
        py_var = utils.fix_python_overlap(var)
        ioc_str += indent_str + "try:\n"
        ioc_str += indent_str + " "*4 + "vm_context.save_intermediate_iocs(" + py_var + ")\n"
        ioc_str += indent_str + "except:\n"
        ioc_str += indent_str + " "* 4 + "pass\n"
    return ioc_str
Example #2
0
    def to_python(self, context, params=None, indent=0):
        
        # Get the global variables read in the function body.
        tmp_context = Context(context=context)
        global_var_info, _ = _get_var_vals(self, tmp_context, global_only=True)
        
        # Set up the initial values for the global variables.
        global_var_init_str = ""
        indent_str = " " * indent
        for global_var in global_var_info.keys():
            val = global_var_info[global_var]
            global_var_init_str += indent_str + str(global_var) + " = " + str(val) + "\n"
        
        # Make a copy of the context so we can mark variables as function
        # arguments.
        tmp_context = Context(context=context)
        for param in self.params:
            tmp_context.set(param.name, "__FUNC_ARG__")

        # Save the name of the current function so we can handle exit function calls.
        tmp_context.curr_func_name = str(self.name)

        # Global variable initialization goes first.
        r = global_var_init_str
        
        # Define the function prototype.
        func_args = "("
        first = True
        for param in self.params:
            if (not first):
                func_args += ", "
            first = False
            func_args += utils.fix_python_overlap(to_python(param, tmp_context))
        func_args += ")"
        r += indent_str + "def " + str(self.name) + func_args + ":\n"

        # Init return value.
        r += indent_str + " " * 4 + "import core.vba_library\n"
        r += indent_str + " " * 4 + "global vm_context\n\n"
        r += indent_str + " " * 4 + "# Function return value.\n"
        r += indent_str + " " * 4 + str(self.name) + " = 0\n\n"

        # Global variables used in the function.
        r += indent_str + " " * 4 + "# Referenced global variables.\n"
        for global_var in global_var_info.keys():
            r += indent_str + " " * 4 + "global " + str(global_var) + "\n"
        r += "\n"
            
        # Function body.
        r += to_python(self.statements, tmp_context, indent=indent+4, statements=True)

        # Check for IOCs.
        r += "\n" + _check_for_iocs(self, tmp_context, indent=indent+4)
        
        # Return the function return val.
        r += "\n" + indent_str + " " * 4 + "return " + str(self.name) + "\n"

        # Done.
        return r
def _updated_vars_to_python(loop, context, indent):
    """Generate Python JIT code for saving the variables updated in a loop
    in Python. These updates are saved in the Python var_updates
    variable.

    @param loop (VBA_Object object) The loop for which to generate
    Python JIT code.

    @param context (Context object) The current program state.

    @param indent (int) The number of spaces to indent the generated
    Python code.

    @return (str) Python JIT code.

    """
    import statements
    
    indent_str = " " * indent
    lhs_visitor = lhs_var_visitor()
    loop.accept(lhs_visitor)
    lhs_var_names = lhs_visitor.variables
    # Handle With variables if needed.
    if (context.with_prefix_raw is not None):
        lhs_var_names.add(safe_str_convert(context.with_prefix_raw))
    # Handle For loop index variables if needed.
    if (isinstance(loop, statements.For_Statement)):
        lhs_var_names.add(safe_str_convert(loop.name))
    var_dict_str = "{"
    first = True
    for var in lhs_var_names:
        py_var = utils.fix_python_overlap(var)
        if (not first):
            var_dict_str += ", "
        first = False
        var = var.replace(".", "")
        var_dict_str += '"' + var + '" : ' + py_var
    var_dict_str += "}"
    save_vals = indent_str + "try:\n"
    save_vals += indent_str + " " * 4 + "var_updates\n"
    save_vals += indent_str + " " * 4 + "var_updates.update(" + var_dict_str + ")\n"
    save_vals += indent_str + "except (NameError, UnboundLocalError):\n"
    save_vals += indent_str + " " * 4 + "var_updates = " + var_dict_str + "\n"
    save_vals += indent_str + 'var_updates["__shell_code__"] = core.vba_library.get_raw_shellcode_data()\n'
    save_vals = indent_str + "# Save the updated variables for reading into ViperMonkey.\n" + save_vals
    if (log.getEffectiveLevel() == logging.DEBUG):
        save_vals += indent_str + "print \"UPDATED VALS!!\"\n"
        save_vals += indent_str + "print var_updates\n"
    return save_vals
def _get_var_vals(item, context, global_only=False):
    """Get the current values for all of the referenced VBA variables
    that appear in the given VBA object.

    @param item (VBA_Object object) The chunk of code to scan to find
    referenced variables.

    @param context (Context object) The current program state.

    @param global_only (boolean) If True only return global variables,
    if False get all variables (local and global).
    
    @return (dict) Returns a dict mapping var names to values.

    """

    import procedures
    import statements

    # Get all the variables.

    # Vars on RHS.
    var_visitor = var_in_expr_visitor(context)
    item.accept(var_visitor, no_embedded_loops=False)
    var_names = var_visitor.variables

    # Vars on LHS.
    lhs_visitor = lhs_var_visitor()
    item.accept(lhs_visitor, no_embedded_loops=False)
    lhs_var_names = lhs_visitor.variables
    
    # Handle member access expressions.
    var_names = var_names.union(lhs_var_names)
    tmp = set()
    for var in var_names:
        tmp.add(var)
        if ("." in var):
            tmp.add(var[:var.index(".")])
    var_names = tmp

    # Handle With variables if needed.
    if (context.with_prefix_raw is not None):
        var_names.add(safe_str_convert(context.with_prefix_raw))
    
    # Get a value for each variable.
    r = {}
    zero_arg_funcs = set()
    for var in var_names:

        # Don't try to convert member access expressions that involve
        # method calls to Python variables. These should be handled
        # later as actual calls.
        if ("(" in var):
            continue

        # Do we already know the variable value?        
        val = None
        orig_val = None
        try:

            # Try to get the current value.
            val = context.get(var, global_only=global_only)
            orig_val = val
            
            # We have been kind of fuzzing the distinction between global and
            # local variables, so tighten down on globals only by just picking
            # up global variables that appear on the RHS but not LHS.
            if (global_only and (var in lhs_var_names)):
                continue
            
            # Do not set function arguments to new values.
            # Do not set loop index variables to new values.
            if ((val == "__FUNC_ARG__") or
                (val == "__ALREADY_SET__") or
                (val == "__LOOP_VAR__")):
                continue
            
            # Function definitions are not valid values.
            if isinstance(val, (VbaLibraryFunc, procedures.Function, procedures.Sub, statements.External_Function)):

                # Don't use the function definition as the value.
                val = None
                
                # 0 arg func calls should only appear on the RHS
                if (var not in lhs_var_names):
                    zero_arg_funcs.add(var)

                    # Don't treat these function calls as variables and
                    # assign initial values to them.
                    context.set("__ORIG__" + var, orig_val, force_local=True)
                    context.set("__ORIG__" + var, orig_val, force_global=True)
                    continue

            # 'inf' is not a valid value.
            val_str = None
            try:
                val_str = safe_str_convert(val).strip()
            except UnicodeEncodeError:
                val_str = filter(isprint, val).strip()
            if ((val_str == "inf") or
                (val_str == "-inf")):
                val = None

            # 'NULL' is not a valid value.
            if (val_str == "NULL"):
                val = None

            # Weird bug.
            if ("core.vba_library.run_function" in val_str):
                val = 0
            
        # Unedfined variable.
        except KeyError:
            if global_only:
                continue

        # Got a valid value for the variable?
        if (val is None):

            # Variable is not defined. Try to infer the type based on how it is used.
            #print "TOP LOOK TYPE: " + safe_str_convert(var)
            var_type, certain_of_type = _infer_type(var, item, context)
            #print (var_type, certain_of_type)
            if (var_type == "INTEGER"):
                val = "NULL"
                if certain_of_type:
                    #print "SET TYPE INT"
                    #print var
                    val = 0
                    context.set_type(var, "Integer")
            elif (var_type == "STRING"):
                val = ""
                if certain_of_type:
                    context.set_type(var, "String")
            else:
                log.warning("Type '" + safe_str_convert(var_type) + "' of var '" + safe_str_convert(var) + "' not handled." + \
                            " Defaulting initial value to \"NULL\".")
                val = "NULL"

        # Rename some vars that overlap with python builtins.
        var = utils.fix_python_overlap(var)
            
        # Save the variable value.
        r[var] = val

        # Save the regex pattern if this is a regex object.
        if (safe_str_convert(val) == "RegExp"):
            if (context.contains("RegExp.pattern")):
                pval = to_python(context.get("RegExp.pattern"), context)
                if (pval.startswith('"')):
                    pval = pval[1:]
                if (pval.endswith('"')):
                    pval = pval[:-1]
                r[var + ".Pattern"] = pval
            if (context.contains("RegExp.global")):
                gval = to_python(context.get("RegExp.global"), context)
                gval = gval.replace('"', "")
                if (gval == "True"):
                    gval = True
                if (gval == "False"):
                    gval = False
                r[var + ".Global"] = gval
        
        # Mark this variable as being set in the Python code to avoid
        # embedded loop Python code generation stomping on the value.
        context.set(var, "__ALREADY_SET__", force_local=True)
        context.set(var, "__ALREADY_SET__", force_global=True)
        
        # Save the original value so we know it's data type for later use in JIT
        # code generation.
        if (orig_val is None):
            orig_val = val
        context.set("__ORIG__" + var, orig_val, force_local=True)
        context.set("__ORIG__" + var, orig_val, force_global=True)
        
    # Done.
    return (r, zero_arg_funcs)