Repository: mandiant/capa
File: tests/test_strings.py (Python)

Code before:
    extract_ascii_strings,
    extract_unicode_strings,
)


def test_buf_filled_with():
    # single repeating byte
    assert buf_filled_with(b"\x00" * 8, 0x00) is True
    assert buf_filled_with(b"\xFF" * 8, 0xFF) is True

    # mixed bytes
    assert buf_filled_with(b"\x00\x01" * 8, 0x00) is False
    assert buf_filled_with(b"ABCD" * 8, ord("A")) is False

    # edge cases
    assert buf_filled_with(b"", 0x00) is False  # Empty buffer
    assert buf_filled_with(b"\x00", 0x00) is True  # Single byte


def test_extract_ascii_strings():
    buf = b"Hello World\x00This is a test\x00"
    strings = list(extract_ascii_strings(buf))
    assert len(strings) == 2
    assert strings[0] == String("Hello World", 0)
    assert strings[1] == String("This is a test", 12)

    # min length
    buf = b"Hi\x00Test\x00"
    strings = list(extract_ascii_strings(buf, n=4))
    assert len(strings) == 1
    assert strings[0] == String("Test", 3)

    # non-ASCII strings
    buf = b"Hello\xFFWorld\x00"
    strings = list(extract_ascii_strings(buf))
    assert len(strings) == 2
    assert strings[0] == String("Hello", 0)
    assert strings[1] == String("World", 6)

    buf = b"\x00" * 8 + b"ValidString\x00"
    strings = list(extract_ascii_strings(buf))
    assert len(strings) == 1
    assert strings[0] == String("ValidString", 8)


def test_extract_unicode_strings():
    buf = b"H\x00e\x00l\x00l\x00o\x00\x00\x00"
    strings = list(extract_unicode_strings(buf))
    assert len(strings) == 1
    assert strings[0] == String("Hello", 0)


Review comment (bug):
Some edge cases to consider for `test_extract_ascii_strings`:
```py
    # test empty buffer
    assert list(extract_ascii_strings(b"")) == []
    
    # Test buffer with only non-ASCII
    assert list(extract_ascii_strings(b"\xFF\xFF\xFF")) == []
```

Code after:
# limitations under the License.

from capa.features.extractors.strings import (
    String,
    buf_filled_with,
    is_printable_str,
    extract_ascii_strings,
    extract_unicode_strings,
)


def test_buf_filled_with():
    # single repeating byte
    assert buf_filled_with(b"\x00" * 8, 0x00) is True
    assert buf_filled_with(b"\xFF" * 8, 0xFF) is True

    # mixed bytes
    assert buf_filled_with(b"\x00\x01" * 8, 0x00) is False
    assert buf_filled_with(b"ABCD" * 8, ord("A")) is False

    # edge cases
    assert buf_filled_with(b"", 0x00) is False  # Empty buffer
    assert buf_filled_with(b"\x00", 0x00) is True  # Single byte


def test_extract_ascii_strings():
    # test empty buffer
    assert list(extract_ascii_strings(b"")) == []

    buf = b"Hello World\x00This is a test\x00"
    strings = list(extract_ascii_strings(buf))
    assert len(strings) == 2
    assert strings[0] == String("Hello World", 0)
    assert strings[1] == String("This is a test", 12)

    # min length
    buf = b"Hi\x00Test\x00"
    strings = list(extract_ascii_strings(buf, n=4))
    assert len(strings) == 1
    assert strings[0] == String("Test", 3)

    # non-ASCII strings
    buf = b"Hello\xFFWorld\x00"
    strings = list(extract_ascii_strings(buf))
    assert len(strings) == 2
    assert strings[0] == String("Hello", 0)
    assert strings[1] == String("World", 6)

    # only non-ASCII
    assert list(extract_ascii_strings(b"\xFF\xFF\xFF")) == []