Repository: mandiant/capa
File: capa/features/extractors/cape/thread.py (Python)

Code before:

def extract_call_features(behavior: Dict, ph: ProcessHandle, th: ThreadHandle) -> Iterator[Tuple[Feature, Address]]:
    """
    this method goes through the specified thread's call trace, and extracts all possible
    features such as: API, Number (for arguments), String (for arguments).

    args:
      behavior: a dictionary of behavioral artifacts extracted by the sandbox
      ph: process handle (for defining the extraction scope)
      th: thread handle (for defining the extraction scope)

    yields:
      Feature, address; where Feature is either: API, Number, or String.
    """

    process = capa.features.extractors.cape.helpers.find_process(behavior["processes"], ph)
    calls: List[Dict[str, Any]] = process["calls"]

    tid = str(th.tid)
    for call in calls:
        if call["thread_id"] != tid:
            continue

        # TODO this address may vary from the PE header, may read actual base from procdump.pe.imagebase or similar
        caller = DynamicAddress(call["id"], int(call["caller"], 16))
        # list similar to disassembly: arguments left-to-right, call
        for arg in call["arguments"][::-1]:
            try:
                yield Number(int(arg["value"], 16), description=f"{arg['name']}"), caller
            except ValueError:
                yield String(arg["value"], description=f"{arg['name']}"), caller
        yield API(call["api"]), caller


def extract_features(behavior: Dict, ph: ProcessHandle, th: ThreadHandle) -> Iterator[Tuple[Feature, Address]]:
    for handler in THREAD_HANDLERS:
        for feature, addr in handler(behavior, ph, th):
            yield feature, addr


THREAD_HANDLERS = (extract_call_features,)


Review comment (suggestion):
```suggestion
        # list similar to disassembly: arguments right-to-left, call
```

Code after:

def extract_call_features(behavior: Dict, ph: ProcessHandle, th: ThreadHandle) -> Iterator[Tuple[Feature, Address]]:
    """
    this method goes through the specified thread's call trace, and extracts all possible
    features such as: API, Number (for arguments), String (for arguments).

    args:
      behavior: a dictionary of behavioral artifacts extracted by the sandbox
      ph: process handle (for defining the extraction scope)
      th: thread handle (for defining the extraction scope)

    yields:
      Feature, address; where Feature is either: API, Number, or String.
    """

    process = capa.features.extractors.cape.helpers.find_process(behavior["processes"], ph)
    calls: List[Dict[str, Any]] = process["calls"]

    tid = str(th.tid)
    for call in calls:
        if call["thread_id"] != tid:
            continue

        # TODO this address may vary from the PE header, may read actual base from procdump.pe.imagebase or similar
        caller = DynamicAddress(call["id"], int(call["caller"], 16))
        # list similar to disassembly: arguments right-to-left, call
        for arg in call["arguments"][::-1]:
            try:
                yield Number(int(arg["value"], 16), description=f"{arg['name']}"), caller
            except ValueError:
                yield String(arg["value"], description=f"{arg['name']}"), caller
        yield API(call["api"]), caller


def extract_features(behavior: Dict, ph: ProcessHandle, th: ThreadHandle) -> Iterator[Tuple[Feature, Address]]:
    for handler in THREAD_HANDLERS:
        for feature, addr in handler(behavior, ph, th):
            yield feature, addr


THREAD_HANDLERS = (extract_call_features,)