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

Code before:
    #
    # this is also how calls to internal functions may be decoded on x32 and x64.
    # see Lab21-01.exe_:0x140001178
    #
    # follow chained thunks, e.g. in 82bf6347acf15e5d883715dc289d8a2b at 0x14005E0FF in
    # 0x140059342 (viv) / 0x14005E0C0 (IDA)
    # 14005E0FF call    j_ElfClearEventLogFileW (14005AAF8)
    #   14005AAF8 jmp     ElfClearEventLogFileW (14005E196)
    #     14005E196 jmp     cs:__imp_ElfClearEventLogFileW

    elif isinstance(insn.opers[0], envi.archs.i386.disasm.i386PcRelOper):
        imports = get_imports(f.vw)
        target = capa.features.extractors.viv.helpers.get_coderef_from(f.vw, insn.va)
        if not target:
            return

        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

        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


Review comment (refactor):
given that the check for "target is a library function" may early return before we get to the symtab extraction, let's move the symtab extraction before the library function check. since the symtab extraction may extract multiple additional features, it would be good to always get those.

Code after:
    #
    # this is also how calls to internal functions may be decoded on x32 and x64.
    # see Lab21-01.exe_:0x140001178
    #
    # follow chained thunks, e.g. in 82bf6347acf15e5d883715dc289d8a2b at 0x14005E0FF in
    # 0x140059342 (viv) / 0x14005E0C0 (IDA)
    # 14005E0FF call    j_ElfClearEventLogFileW (14005AAF8)
    #   14005AAF8 jmp     ElfClearEventLogFileW (14005E196)
    #     14005E196 jmp     cs:__imp_ElfClearEventLogFileW

    elif isinstance(insn.opers[0], envi.archs.i386.disasm.i386PcRelOper):
        imports = get_imports(f.vw)
        target = capa.features.extractors.viv.helpers.get_coderef_from(f.vw, insn.va)
        if not target:
            return

        if f.vw.metadata["Format"] == "elf":
            if "symtab" not in fh.ctx["cache"]:
                # 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):