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

Code before:
import capa.features.extractors.common
import capa.features.extractors.viv.file
import capa.features.extractors.viv.insn
import capa.features.extractors.viv.global_
import capa.features.extractors.viv.function
import capa.features.extractors.viv.basicblock
from capa.features.common import Feature
from capa.features.address import Address, AbsoluteVirtualAddress
from capa.features.extractors.base_extractor import BBHandle, InsnHandle, FunctionHandle, FeatureExtractor

logger = logging.getLogger(__name__)


class VivisectFeatureExtractor(FeatureExtractor):
    def __init__(self, vw, path, os):
        super().__init__()
        self.vw = vw
        self.path = path
        self.os = os
        with open(self.path, "rb") as f:
            self.buf = f.read()

        # pre-compute these because we'll yield them at *every* scope.
        self.global_features: List[Tuple[Feature, Address]] = []
        self.global_features.extend(capa.features.extractors.viv.file.extract_file_format(self.buf))
        self.global_features.extend(capa.features.extractors.common.extract_os(self.buf, self.os))
        self.global_features.extend(capa.features.extractors.viv.global_.extract_arch(self.vw))

    def get_base_address(self):
        # assume there is only one file loaded into the vw
        return AbsoluteVirtualAddress(list(self.vw.filemeta.values())[0]["imagebase"])

    def extract_global_features(self):
        yield from self.global_features

    def extract_file_features(self):
        yield from capa.features.extractors.viv.file.extract_features(self.vw, self.buf)

    def get_functions(self) -> Iterator[FunctionHandle]:
        for va in sorted(self.vw.getFunctions()):
            yield FunctionHandle(address=AbsoluteVirtualAddress(va), inner=viv_utils.Function(self.vw, va))

    def extract_function_features(self, fh: FunctionHandle) -> Iterator[Tuple[Feature, Address]]:
        yield from capa.features.extractors.viv.function.extract_features(fh)

    def get_basic_blocks(self, fh: FunctionHandle) -> Iterator[BBHandle]:
        f: viv_utils.Function = fh.inner
        for bb in f.basic_blocks:
            yield BBHandle(address=AbsoluteVirtualAddress(bb.va), inner=bb)

    def extract_basic_block_features(self, fh: FunctionHandle, bbh) -> Iterator[Tuple[Feature, Address]]:

Review comment (refactor):
```suggestion
        self.global_features.extend(capa.features.extractors.common.extract_os(self.buf, os))
```

Code after:
import capa.features.extractors.common
import capa.features.extractors.viv.file
import capa.features.extractors.viv.insn
import capa.features.extractors.viv.global_
import capa.features.extractors.viv.function
import capa.features.extractors.viv.basicblock
from capa.features.common import Feature
from capa.features.address import Address, AbsoluteVirtualAddress
from capa.features.extractors.base_extractor import BBHandle, InsnHandle, FunctionHandle, FeatureExtractor

logger = logging.getLogger(__name__)


class VivisectFeatureExtractor(FeatureExtractor):
    def __init__(self, vw, path, os):
        super().__init__()
        self.vw = vw
        self.path = path
        with open(self.path, "rb") as f:
            self.buf = f.read()

        # pre-compute these because we'll yield them at *every* scope.
        self.global_features: List[Tuple[Feature, Address]] = []
        self.global_features.extend(capa.features.extractors.viv.file.extract_file_format(self.buf))
        self.global_features.extend(capa.features.extractors.common.extract_os(self.buf, os))
        self.global_features.extend(capa.features.extractors.viv.global_.extract_arch(self.vw))

    def get_base_address(self):
        # assume there is only one file loaded into the vw
        return AbsoluteVirtualAddress(list(self.vw.filemeta.values())[0]["imagebase"])

    def extract_global_features(self):
        yield from self.global_features

    def extract_file_features(self):
        yield from capa.features.extractors.viv.file.extract_features(self.vw, self.buf)

    def get_functions(self) -> Iterator[FunctionHandle]:
        for va in sorted(self.vw.getFunctions()):
            yield FunctionHandle(address=AbsoluteVirtualAddress(va), inner=viv_utils.Function(self.vw, va))

    def extract_function_features(self, fh: FunctionHandle) -> Iterator[Tuple[Feature, Address]]:
        yield from capa.features.extractors.viv.function.extract_features(fh)

    def get_basic_blocks(self, fh: FunctionHandle) -> Iterator[BBHandle]:
        f: viv_utils.Function = fh.inner
        for bb in f.basic_blocks:
            yield BBHandle(address=AbsoluteVirtualAddress(bb.va), inner=bb)

    def extract_basic_block_features(self, fh: FunctionHandle, bbh) -> Iterator[Tuple[Feature, Address]]:
        yield from capa.features.extractors.viv.basicblock.extract_features(fh, bbh)