From bd5c65d22c2de0b3f396fc37b0726a5ce8b69db7 Mon Sep 17 00:00:00 2001
From: Willi Ballenthin <wballenthin@google.com>
Date: Mon, 14 Aug 2023 11:08:19 +0000
Subject: [PATCH] elf: fix parsing of symtab from viv

closes #1704
---
 capa/features/extractors/elf.py          | 18 ++++++++++--------
 capa/features/extractors/viv/function.py |  2 +-
 capa/features/extractors/viv/insn.py     |  2 +-
 3 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/capa/features/extractors/elf.py b/capa/features/extractors/elf.py
index 01662a124..62846ea85 100644
--- a/capa/features/extractors/elf.py
+++ b/capa/features/extractors/elf.py
@@ -13,6 +13,8 @@
 from typing import Set, Dict, List, Tuple, BinaryIO, Iterator, Optional
 from dataclasses import dataclass
 
+import Elf  # from vivisect
+
 logger = logging.getLogger(__name__)
 
 
@@ -710,17 +712,17 @@ def get_symbols(self) -> Iterator[Symbol]:
         yield from self.symbols
 
     @classmethod
-    def from_Elf(cls, ElfBinary) -> Optional["SymTab"]:
-        endian = "<" if ElfBinary.getEndian() == 0 else ">"
-        bitness = ElfBinary.bits
+    def from_viv(cls, elf: Elf.Elf) -> Optional["SymTab"]:
+        endian = "<" if elf.getEndian() == 0 else ">"
+        bitness = elf.bits
 
         SHT_SYMTAB = 0x2
-        for section in ElfBinary.sections:
-            if section.sh_info & SHT_SYMTAB:
-                strtab_section = ElfBinary.sections[section.sh_link]
-                sh_symtab = Shdr.from_viv(section, ElfBinary.readAtOffset(section.sh_offset, section.sh_size))
+        for section in elf.sections:
+            if section.sh_type & SHT_SYMTAB:
+                strtab_section = elf.sections[section.sh_link]
+                sh_symtab = Shdr.from_viv(section, elf.readAtOffset(section.sh_offset, section.sh_size))
                 sh_strtab = Shdr.from_viv(
-                    strtab_section, ElfBinary.readAtOffset(strtab_section.sh_offset, strtab_section.sh_size)
+                    strtab_section, elf.readAtOffset(strtab_section.sh_offset, strtab_section.sh_size)
                 )
 
         try:
diff --git a/capa/features/extractors/viv/function.py b/capa/features/extractors/viv/function.py
index 112f4fa18..b018b34bb 100644
--- a/capa/features/extractors/viv/function.py
+++ b/capa/features/extractors/viv/function.py
@@ -38,7 +38,7 @@ def extract_function_symtab_names(fh: FunctionHandle) -> Iterator[Tuple[Feature,
         # this is in order to eliminate the computational overhead of refetching symtab each time.
         if "symtab" not in fh.ctx["cache"]:
             try:
-                fh.ctx["cache"]["symtab"] = SymTab.from_Elf(fh.inner.vw.parsedbin)
+                fh.ctx["cache"]["symtab"] = SymTab.from_viv(fh.inner.vw.parsedbin)
             except Exception:
                 fh.ctx["cache"]["symtab"] = None
 
diff --git a/capa/features/extractors/viv/insn.py b/capa/features/extractors/viv/insn.py
index fb117dfa3..7b88dd2de 100644
--- a/capa/features/extractors/viv/insn.py
+++ b/capa/features/extractors/viv/insn.py
@@ -115,7 +115,7 @@ def extract_insn_api_features(fh: FunctionHandle, bb, ih: InsnHandle) -> Iterato
                 # 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)
+                    fh.ctx["cache"]["symtab"] = SymTab.from_viv(f.vw.parsedbin)
                 except Exception:
                     fh.ctx["cache"]["symtab"] = None