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

Code before:
            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

        if f.vw.metadata["Format"] == "elf":
            if "SymbolTable" not in f.vw.metadata:
                # 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.vsGetField("sh_info") & SHT_SYMTAB != 0:
                        strtab = elf.sections[section.vsGetField("sh_link")]
                        sh_symtab = Shdr.from_viv(section, elf.getSectionBytes(section.name))
                        sh_strtab = Shdr.from_viv(strtab, elf.getSectionBytes(strtab.name))

                symtab = SymTab(endian, bitness, sh_symtab, sh_strtab)
                f.vw.metadata["SymbolTable"] = symtab

            symtab = f.vw.metadata["SymbolTable"]
            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), ih.address

        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):
i don't think we should store the symtab cache in the vivisect workspace for a couple reasons:
  1. we don't really own this data structure and it could change (admittedly unlikely), and
  2. the cache probably isn't serializable, and I'm not sure how this works (or doesn't) with saving the viv workspace.

for (2), I'm thinking that a `SymTab` is a python-native object, so it can't be serialized to msgpack or json (i think what viv uses to save the .viv file). so, this may raise an exception during saving, or maybe not write it at all, and in either case, I'm not sure what happens when it gets reloaded back from disk. all in all, i have a lot of questions about how this works, and it would be simpler to put this in a place we control.

i recommend that we use the `FunctionHandle.ctx` field to store our cached `SymTab` instance.

we should update here: https://github.com/mandiant/capa/blob/d3c44a82633ca2ee5cae2182197dc3706bd7c0be/capa/features/extractors/viv/extractor.py#L53

to pass in a value for the `ctx=` kwarg declared here: https://github.com/mandiant/capa/blob/d3c44a82633ca2ee5cae2182197dc3706bd7c0be/capa/features/extractors/base_extractor.py#L37

maybe something like:

```python
cache = {}
for va in ...:
  yield FunctionHandle(..., ctx={"cache": cache})
```

this way we have a single cache instance that's shared across all the function handles.

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: