"""Entity selectors for resolving function, address, and entity references. Selectors are human-readable strings that resolve to specific entities. Supported selector formats: - function: — Resolve a function by name (exact or fuzzy match) - function:
— Resolve a function by address - address: — Resolve an address range (e.g., 0x1000..0x2000) - name: — Generic entity lookup by name """ from __future__ import annotations import re from dataclasses import dataclass, field from typing import Any from binary_analysis.domain.entities import Function from binary_analysis.domain.errors import AmbiguousSelectorError, EntityNotFoundError # --------------------------------------------------------------------------- # Selector types # --------------------------------------------------------------------------- class SelectorKind: """Selector kind constants.""" FUNCTION = "function" ADDRESS = "address" NAME = "name" # --------------------------------------------------------------------------- # Parsed selector # --------------------------------------------------------------------------- @dataclass class ParsedSelector: """Result of parsing an entity selector string. Attributes: kind: The selector kind (function, address, name). value: The parsed selector value. raw: The original selector string. is_address: Whether the value represents an address. address_value: Parsed address offset (hex string without 0x prefix) if applicable. is_range: Whether the selector specifies a range. range_start: Start of range if is_range is True. range_end: End of range if is_range is True. """ kind: str = "" value: str = "" raw: str = "" is_address: bool = False address_value: str | None = None is_range: bool = False range_start: str | None = None range_end: str | None = None def __str__(self) -> str: return self.raw # --------------------------------------------------------------------------- # Selector parser # --------------------------------------------------------------------------- _ADDRESS_PATTERN = re.compile(r"^(0x)?[0-9a-fA-F]+$") _RANGE_PATTERN = re.compile(r"^(0x)?[0-9a-fA-F]+\.\.(0x)?[0-9a-fA-F]+$") _SELECTOR_PATTERN = re.compile(r"^(function|address|name):(.+)$", re.IGNORECASE) def parse_selector(raw: str) -> ParsedSelector: """Parse a selector string into its components. Supported formats: "function:main" → function selector by name "function:0x401000" → function selector by address "address:0x1000..0x2000" → address range selector "name:entrypoint" → generic name selector "main" → implicit function selector (shorthand) "0x401000" → implicit address selector (shorthand) Args: raw: The raw selector string. Returns: A ParsedSelector with kind, value, and parsed components. """ result = ParsedSelector(raw=raw, kind=SelectorKind.NAME, value=raw) # Try explicit selector format: kind:value match = _SELECTOR_PATTERN.match(raw) if match: kind = match.group(1).lower() value = match.group(2) if kind == SelectorKind.FUNCTION: result.kind = SelectorKind.FUNCTION result.value = value elif kind == SelectorKind.ADDRESS: result.kind = SelectorKind.ADDRESS result.value = value else: result.kind = SelectorKind.NAME result.value = value else: # Implicit: check if it looks like an address if _ADDRESS_PATTERN.match(raw): result.kind = SelectorKind.ADDRESS result.value = raw else: result.kind = SelectorKind.FUNCTION result.value = raw # Check if it's an address value if _ADDRESS_PATTERN.match(result.value): result.is_address = True addr = result.value if addr.startswith("0x") or addr.startswith("0X"): result.address_value = addr[2:].lower() else: result.address_value = addr.lower() # Check if it's a range if _RANGE_PATTERN.match(result.value): result.is_range = True parts = result.value.split("..") result.range_start = parts[0] result.range_end = parts[1] return result # --------------------------------------------------------------------------- # Entity resolver # --------------------------------------------------------------------------- @dataclass class ResolvedEntity: """Result of resolving a selector to one or more entities. Attributes: selector: The parsed selector that was resolved. entity_type: The type of entity resolved (e.g., "Function", "Address"). exact_match: The single entity if resolution was unambiguous. candidates: List of candidates if multiple matches were found. is_ambiguous: Whether resolution produced multiple candidates. """ selector: ParsedSelector entity_type: str = "" exact_match: Any | None = None candidates: list[Any] = field(default_factory=list) is_ambiguous: bool = False def resolve_function( parsed: ParsedSelector, functions: list[Function], require_unique: bool = True, ) -> Function: """Resolve a function selector to a single Function entity. Args: parsed: The parsed function selector. functions: List of functions to search. require_unique: If True, raise AmbiguousSelectorError when multiple functions match. Returns: The matching Function entity. Raises: EntityNotFoundError: If no function matches the selector. AmbiguousSelectorError: If multiple functions match and require_unique is True. """ if parsed.is_address: # Lookup by address addr_val = parsed.address_value or "" matches = [ f for f in functions if f.address is not None and f.address.offset.lower() == f"0x{addr_val}" ] if not matches: matches = [ f for f in functions if f.address is not None and addr_val in f.address.offset.lower() ] else: # Lookup by name search_name = parsed.value.lower() exact_matches = [f for f in functions if f.name.lower() == search_name] matches = exact_matches or [f for f in functions if search_name in f.name.lower()] if not matches: raise EntityNotFoundError("Function", parsed.raw) if len(matches) > 1 and require_unique: candidates_info = [ { "name": f.name, "address": f.address.to_dict() if f.address else None, "size_bytes": f.size_bytes, } for f in matches ] raise AmbiguousSelectorError( f"Function selector '{parsed.raw}' matches {len(matches)} functions", candidates=candidates_info, ) return matches[0] def resolve_functions( parsed: ParsedSelector, functions: list[Function], ) -> list[Function]: """Resolve a function selector to all matching Function entities. Args: parsed: The parsed function selector. functions: List of functions to search. Returns: List of matching Function entities (may be empty). """ if parsed.is_address: addr_val = parsed.address_value or "" matches = [ f for f in functions if f.address is not None and addr_val in f.address.offset.lower() ] else: search_name = parsed.value.lower() matches = [f for f in functions if search_name in f.name.lower()] return matches def format_candidates(candidates: list[dict[str, Any]]) -> str: """Format candidate entities for display in ambiguity errors. Args: candidates: List of candidate dicts with name, address, and optional info. Returns: A human-readable string listing candidates. """ lines = ["Ambiguous selector matches multiple entities:"] for i, candidate in enumerate(candidates, start=1): name = candidate.get("name", "unknown") addr = candidate.get("address", {}) if isinstance(addr, dict): display = addr.get("display", addr.get("offset", "?")) else: display = str(addr) lines.append(f" {i}. {name} @ {display}") return "\n".join(lines)