Repository: mandiant/capa
File: capa/features/extractors/viv/insn.py (Python)

Code before:
                # the symbol table gets stored as a function's attribute in order to avoid running
                # this code everytime the call is made, thus preventing the computational overhead.
                elf = f.vw.parsedbin
                endian = "<" if elf.getEndian() == 0 else ">"
                bitness = elf.bits

                SHT_SYMTAB = 0x2
                for section in elf.sections:
                    if section.sh_info & SHT_SYMTAB:
                        strtab_section = elf.sections[section.vsGetField("sh_link")]
                        sh_symtab = Shdr.from_viv(section, elf.readAtOffset(section.sh_offset, section.sh_size))
                        sh_strtab = Shdr.from_viv(strtab_section, elf.readAtOffset(strtab_section.sh_offset, strtab_section.sh_size))

                symtab = SymTab(endian, bitness, sh_symtab, sh_strtab)
                fh.ctx["cache"]["symtab"] = symtab

            symtab = fh.ctx["cache"]["symtab"]
            for symbol in symtab.get_symbols():
                sym_name = symtab.get_name(symbol)
                sym_value = symbol.value
                sym_info = symbol.info

                STT_FUNC = 0x2
                if sym_value == target and sym_info & STT_FUNC != 0:
                    yield API(sym_name), ih.address
                    yield FunctionName(sym_name), target

        if viv_utils.flirt.is_library_function(f.vw, target):
            name = viv_utils.get_function_name(f.vw, target)
            yield API(name), ih.address
            if name.startswith("_"):
                # some linkers may prefix linked routines with a `_` to avoid name collisions.
                # extract features for both the mangled and un-mangled representations.
                # e.g. `_fwrite` -> `fwrite`
                # see: https://stackoverflow.com/a/2628384/87207
                yield API(name[1:]), ih.address
            return

        for _ in range(THUNK_CHAIN_DEPTH_DELTA):
            if target in imports:
                dll, symbol = imports[target]
                for name in capa.features.extractors.helpers.generate_symbols(dll, symbol):
                    yield API(name), ih.address

            # if jump leads to an ENDBRANCH instruction, skip it
            if f.vw.getByteDef(target)[1].startswith(b"\xf3\x0f\x1e"):
                target += 4

            target = capa.features.extractors.viv.helpers.get_coderef_from(f.vw, target)
            if not target:
                return

Review comment (performance):
this should be extracted at function scope, for the target function. extracting it here sorta works, but its not correct.

imagine the case where we might ask capa: "show me all features for the function at 0x401000" and it extracts features only for that function. we'd want the `FunctionName` to be extracted for that function in that function, not when we encounter a call to that function.

so, please move this extraction into `viv/function.py`. this might also require parsing the symtab there, so it might make sense to refactor the symtab cache creation into a function that can be reused in both places.

Code after:
                # the symbol table gets stored as a function's attribute in order to avoid running
                # this code everytime the call is made, thus preventing the computational overhead.
                try:
                    fh.ctx["cache"]["symtab"] = SymTab.from_Elf(f.vw.parsedbin)
                except:
                    fh.ctx["cache"]["symtab"] = None

            symtab = fh.ctx["cache"]["symtab"]
            if symtab:
                for symbol in symtab.get_symbols():
                    sym_name = symtab.get_name(symbol)
                    sym_value = symbol.value
                    sym_info = symbol.info

                    STT_FUNC = 0x2
                    if sym_value == target and sym_info & STT_FUNC != 0:
                        yield API(sym_name), ih.address

        if viv_utils.flirt.is_library_function(f.vw, target):
            name = viv_utils.get_function_name(f.vw, target)
            yield API(name), ih.address
            if name.startswith("_"):
                # some linkers may prefix linked routines with a `_` to avoid name collisions.
                # extract features for both the mangled and un-mangled representations.
                # e.g. `_fwrite` -> `fwrite`
                # see: https://stackoverflow.com/a/2628384/87207
                yield API(name[1:]), ih.address
            return

        for _ in range(THUNK_CHAIN_DEPTH_DELTA):
            if target in imports:
                dll, symbol = imports[target]
                for name in capa.features.extractors.helpers.generate_symbols(dll, symbol):
                    yield API(name), ih.address

            # if jump leads to an ENDBRANCH instruction, skip it
            if f.vw.getByteDef(target)[1].startswith(b"\xf3\x0f\x1e"):
                target += 4

            target = capa.features.extractors.viv.helpers.get_coderef_from(f.vw, target)
            if not target:
                return

    # call via import on x64
    # see Lab21-01.exe_:0x14000118C
    elif isinstance(insn.opers[0], envi.archs.amd64.disasm.Amd64RipRelOper):
        op = insn.opers[0]
        target = op.getOperAddr(insn)

        imports = get_imports(f.vw)
        if target in imports: