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

Code before:
MATCH_RESULT = b'{"meta":'
MATCH_JSON_OBJECT = b'{"'


def extract_file_strings(buf, **kwargs) -> Iterator[Tuple[String, Address]]:
    """
    extract ASCII and UTF-16 LE strings from file
    """
    for s in capa.features.extractors.strings.extract_ascii_strings(buf):
        yield String(s.s), FileOffsetAddress(s.offset)

    for s in capa.features.extractors.strings.extract_unicode_strings(buf):
        yield String(s.s), FileOffsetAddress(s.offset)


def extract_format(buf) -> Iterator[Tuple[Feature, Address]]:
    if buf.startswith(MATCH_PE):
        yield Format(FORMAT_PE), NO_ADDRESS
    elif buf.startswith(MATCH_ELF):
        yield Format(FORMAT_ELF), NO_ADDRESS
    elif is_freeze(buf):
        yield Format(FORMAT_FREEZE), NO_ADDRESS
    elif buf.startswith(MATCH_RESULT):
        yield Format(FORMAT_RESULT), NO_ADDRESS
    elif (
        buf.replace(b" ", b"").replace(b"\r", b"").replace(b"\n", b"").replace(b"\t", b"").startswith(MATCH_JSON_OBJECT)
    ):
        # potential start of JSON object data - `{"` without whitespace (\r\n\t)
        # we don't know what it is exactly, but may support it (e.g. a dynamic CAPE sandbox report)
        # skip verdict here and let subsequent code analyze this further
        return
    else:
        # we likely end up here:
        #  1. handling a file format (e.g. macho)
        #
        # for (1), this logic will need to be updated as the format is implemented.
        logger.debug("unsupported file format: %s", binascii.hexlify(buf[:4]).decode("ascii"))
        return


def extract_arch(buf) -> Iterator[Tuple[Feature, Address]]:
    if buf.startswith(MATCH_PE):
        yield from capa.features.extractors.pefile.extract_file_arch(pe=pefile.PE(data=buf))

    elif buf.startswith(MATCH_RESULT):
        yield Arch(ARCH_ANY), NO_ADDRESS

    elif buf.startswith(MATCH_ELF):
        with contextlib.closing(io.BytesIO(buf)) as f:
            arch = capa.features.extractors.elf.detect_elf_arch(f)


Review comment (suggestion):
some options:
  - maybe consider only taking a taste of the file, like the first 100 bytes, to avoid all the work of stripping whitespace from parts of the file we won't use here (assume there aren't more than 100 newlines at the start of the file ;-) ).
  - `re.sub("\w", "")` probably replaces all whitespace in one go, rather than the intermediate passes, and is maybe more readable

which would result in something like: `elif re.sub(buf[:100], "\w", "").startswith(MATCH_JSON_OBJECT)`

Code after:
MATCH_ELF = b"\x7fELF"
MATCH_RESULT = b'{"meta":'
MATCH_JSON_OBJECT = b'{"'


def extract_file_strings(buf, **kwargs) -> Iterator[Tuple[String, Address]]:
    """
    extract ASCII and UTF-16 LE strings from file
    """
    for s in capa.features.extractors.strings.extract_ascii_strings(buf):
        yield String(s.s), FileOffsetAddress(s.offset)

    for s in capa.features.extractors.strings.extract_unicode_strings(buf):
        yield String(s.s), FileOffsetAddress(s.offset)


def extract_format(buf) -> Iterator[Tuple[Feature, Address]]:
    if buf.startswith(MATCH_PE):
        yield Format(FORMAT_PE), NO_ADDRESS
    elif buf.startswith(MATCH_ELF):
        yield Format(FORMAT_ELF), NO_ADDRESS
    elif is_freeze(buf):
        yield Format(FORMAT_FREEZE), NO_ADDRESS
    elif buf.startswith(MATCH_RESULT):
        yield Format(FORMAT_RESULT), NO_ADDRESS
    elif re.sub(rb"\w", b"", buf[:20]).startswith(MATCH_JSON_OBJECT):
        # potential start of JSON object data without whitespace
        # we don't know what it is exactly, but may support it (e.g. a dynamic CAPE sandbox report)
        # skip verdict here and let subsequent code analyze this further
        return
    else:
        # we likely end up here:
        #  1. handling a file format (e.g. macho)
        #
        # for (1), this logic will need to be updated as the format is implemented.
        logger.debug("unsupported file format: %s", binascii.hexlify(buf[:4]).decode("ascii"))
        return


def extract_arch(buf) -> Iterator[Tuple[Feature, Address]]:
    if buf.startswith(MATCH_PE):
        yield from capa.features.extractors.pefile.extract_file_arch(pe=pefile.PE(data=buf))

    elif buf.startswith(MATCH_RESULT):
        yield Arch(ARCH_ANY), NO_ADDRESS

    elif buf.startswith(MATCH_ELF):
        with contextlib.closing(io.BytesIO(buf)) as f:
            arch = capa.features.extractors.elf.detect_elf_arch(f)

        if arch not in capa.features.common.VALID_ARCH: