refactor: move layout string extraction into Layout.extract_strings

This commit is contained in:
lakshit verma 2026-08-12 21:40:44 +05:30
parent 35432977f9
commit f62f41d100
5 changed files with 71 additions and 71 deletions

View file

@ -128,6 +128,71 @@ class Layout(BaseModel, abc.ABC):
"convenience"
return self.slice.range.end
def extract_strings(self, min_len: int) -> None:
"""
find the strings in this layout and its children, recursively.
this finds strings in the gaps between the children (and before the
first and after the last child), so this method must run before
``tag_strings``.
"""
# imported here to avoid a circular import with floss.layout.extract
from floss.layout.extract import extract_strings as extract_gap_strings
if not self.children:
# at this moment, self.strings contains only ExtractedStrings
# after tag_strings, it will contain TaggedStrings.
self.strings = extract_gap_strings(self.slice, min_len) # type: ignore
return
# we have children, so we need to recurse to find their strings,
# and also find strings in the gaps between children.
# lets find the gap strings first:
for i, child in enumerate(self.children):
if i == 0:
# find the strings before the first child
offset = 0
size = self.children[0].offset - self.offset
else:
# find strings between children
prior = self.children[i - 1]
offset = prior.end - self.offset
size = child.offset - prior.end
if size == 0:
# there is no gap here.
continue
gap = self.slice.slice(offset, size)
self.strings.extend(extract_gap_strings(gap, min_len)) # type: ignore
# finally, find strings after the last child
last_child = self.children[-1]
offset = last_child.end - self.offset
size = self.end - last_child.end
if size > 0:
gap = self.slice.slice(offset, size)
self.strings.extend(extract_gap_strings(gap, min_len)) # type: ignore
# now recurse to find the strings in the children.
for child in self.children:
child.extract_strings(min_len)
if self.strings:
child_ranges = [(child.offset, child.end) for child in self.children]
filtered = []
for string in self.strings:
if isinstance(string, TaggedString):
offset = string.offset
else:
offset = string.slice.range.offset
if any(start <= offset < end for start, end in child_ranges):
continue
filtered.append(string)
self.strings = filtered
def tag_strings(self, taggers: Sequence[Tagger]):
"""
tag the strings in this layout and its children, recursively.

View file

@ -68,70 +68,6 @@ def extract_strings(slice: Slice, n: int = MIN_STR_LEN) -> Iterable[ExtractedStr
)
def extract_layout_strings(layout: Layout, min_len: int):
if not layout.children:
# all the strings are found in this slice directly.
# at this moment, layout.strings contains only ExtractedStrings
# after layout.tag_strings, it will contain TaggedStrings.
layout.strings = extract_strings(layout.slice, min_len) # type: ignore
return
else:
# we have children, so we need to recurse to find their strings,
# and also find strings in the gaps between children.
# lets find the gap strings first:
for i, child in enumerate(layout.children):
if i == 0:
# find the strings before the first child
offset = 0
size = layout.children[0].offset - layout.offset
else:
# find strings between children
prior = layout.children[i - 1]
offset = prior.end - layout.offset
size = child.offset - prior.end
if size == 0:
# there is no gap here.
continue
gap = layout.slice.slice(offset, size)
# at this moment, layout.strings contains only ExtractedStrings
# after layout.tag_strings, it will contain TaggedStrings.
layout.strings.extend(extract_strings(gap, min_len)) # type: ignore
# finally, find strings after the last child
last_child = layout.children[-1]
offset = last_child.end - layout.offset
size = layout.end - last_child.end
if size > 0:
gap = layout.slice.slice(offset, size)
# at this moment, layout.strings contains only ExtractedStrings
# after layout.tag_strings, it will contain TaggedStrings.
layout.strings.extend(extract_strings(gap, min_len)) # type: ignore
# now recurse to find the strings in the children.
for child in layout.children:
extract_layout_strings(child, min_len)
if layout.strings:
child_ranges = [(child.offset, child.end) for child in layout.children]
filtered = []
for string in layout.strings:
if isinstance(string, TaggedString):
offset = string.offset
else:
offset = string.slice.range.offset
if any(start <= offset < end for start, end in child_ranges):
continue
filtered.append(string)
layout.strings = filtered
def collect_strings(layout: Layout) -> List[TaggedString]:
ret = []

View file

@ -216,7 +216,6 @@ def compute_layout(
"""
from floss.layout import compute_layout as layout_compute
from floss.ranges import Slice
from floss.layout.extract import extract_layout_strings
try:
file_slice = Slice.from_bytes(buf=buf)
@ -226,7 +225,7 @@ def compute_layout(
logger.debug("no structured layout (got %r); using classic static strings", parsed_layout.name)
return None
extract_layout_strings(parsed_layout, min_length)
parsed_layout.extract_strings(min_length)
return parsed_layout
except Exception as e:
logger.warning("layout-aware static analysis failed; using classic statics: %s", e)

View file

@ -24,7 +24,7 @@ from floss.enrich import static_strings_from_layout
from floss.layout import compute_layout
from floss.ranges import Slice
from floss.results import Strings, Analysis, Metadata, ResultLayout, ResultDocument
from floss.layout.extract import collect_strings, extract_layout_strings
from floss.layout.extract import collect_strings
CD = Path(__file__).resolve().parent
MIN_STR_LEN = 6
@ -40,7 +40,7 @@ def analyzed_layout(pma_binary_path):
slice_buf = pma_binary_path.read_bytes()
file_slice = Slice.from_bytes(slice_buf)
parsed = compute_layout(file_slice)
extract_layout_strings(parsed, 6)
parsed.extract_strings(6)
taggers = load_databases()
parsed.tag_strings(taggers)
parsed.mark_structures()
@ -112,7 +112,7 @@ def test_analysis_pipeline(pma_binary_path):
slice_buf = pma_binary_path.read_bytes()
file_slice = Slice.from_bytes(slice_buf)
parsed = compute_layout(file_slice)
extract_layout_strings(parsed, 6)
parsed.extract_strings(6)
# Check that the layout has been computed correctly
assert parsed.name == "pe"

View file

@ -19,7 +19,7 @@ import pytest
from floss.tags import load_databases
from floss.layout import compute_layout
from floss.ranges import Slice
from floss.layout.extract import collect_strings, extract_layout_strings
from floss.layout.extract import collect_strings
@pytest.fixture(scope="module")
@ -33,7 +33,7 @@ def pma0101_layout():
slice_buf = binary_path.read_bytes()
file_slice = Slice.from_bytes(slice_buf)
parsed = compute_layout(file_slice)
extract_layout_strings(parsed, 6)
parsed.extract_strings(6)
taggers = load_databases()
parsed.tag_strings(taggers)
parsed.mark_structures()