policy.py

"""
Core dataclasses for canonical policy representation.

These dataclasses define the normalized forms used for equivalence comparison
across the compile-reverse pipeline.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any
import re

from pawl.normalize.operations import WILDCARD_CHILDREN


@dataclass(frozen=True, eq=False)
class CanonicalPredicate:
    """A normalized filter predicate.

    Predicates are normalized to a canonical string form for comparison.
    The `reconstructed` flag indicates predicates recovered from disconnected
    filter nodes rather than directly reversed from the graph.

    For equivalence purposes, only `text` and `filter_type` are compared.
    The `reconstructed` flag is metadata for provenance tracking.
    """

    text: str  # e.g. '(path-regex #"^/tmp/.*")'
    filter_type: str  # e.g. "path-regex", "iokit-user-client-class"
    reconstructed: bool = False

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, CanonicalPredicate):
            return NotImplemented
        # Only compare text and filter_type, not reconstructed flag
        return self.text == other.text and self.filter_type == other.filter_type

    def __hash__(self) -> int:
        # Hash must be consistent with __eq__
        return hash((self.text, self.filter_type))

    def __str__(self) -> str:
        return self.text


@dataclass(frozen=True)
class CanonicalRule:
    """One normalized sandbox rule.

    Rules are normalized by:
    - Sorting predicates alphabetically
    - Normalizing predicate text (whitespace, quote style)
    - Tracking reconstructed predicates separately
    """

    action: str  # "allow" | "deny"
    operation: str  # e.g. "file-read*", "mach-lookup"
    predicates: frozenset[CanonicalPredicate] = field(default_factory=frozenset)

    @property
    def leaf_predicates(self) -> frozenset[str]:
        """Extract leaf predicate texts, flattening require-any/require-all.

        This enables set-based comparison that ignores grouping differences.
        Source: (allow op (pred-a) (pred-b))
        Reversed: (allow op (require-any (pred-a) (pred-b)))
        Both yield the same leaf set: {pred-a, pred-b}
        """
        leaves: set[str] = set()
        for pred in self.predicates:
            leaves.update(_flatten_predicate_to_leaves(pred.text))
        return frozenset(leaves)

    def semantic_eq(self, other: CanonicalRule) -> bool:
        """Compare rules using flattened leaf predicate sets.

        This is more permissive than __eq__ - it considers rules equivalent
        if they have the same action, operation, and set of leaf predicates,
        even if the grouping structure differs.
        """
        if self.action != other.action or self.operation != other.operation:
            return False
        return self.leaf_predicates == other.leaf_predicates

    def __str__(self) -> str:
        if not self.predicates:
            return f"({self.action} {self.operation})"
        pred_strs = sorted(str(p) for p in self.predicates)
        return f"({self.action} {self.operation}\n\t" + "\n\t".join(pred_strs) + ")"

    @property
    def reconstructed_predicates(self) -> frozenset[CanonicalPredicate]:
        """Predicates that were reconstructed from disconnected filters."""
        return frozenset(p for p in self.predicates if p.reconstructed)

    @property
    def direct_predicates(self) -> frozenset[CanonicalPredicate]:
        """Predicates that were directly reversed from the graph."""
        return frozenset(p for p in self.predicates if not p.reconstructed)


@dataclass
class CanonicalPolicy:
    """Normalized representation of a sandbox policy.

    This is the canonical form used for equivalence comparison.  Two policies
    are equivalent if their CanonicalPolicy representations are equal.

    Attributes
    ----------
    default_action : str
        The default action ("allow" or "deny").
    rules : frozenset[CanonicalRule]
        The normalized rules.  Rules are compared as a set, not ordered.
    metadata : dict
        Provenance information, warnings, and diagnostics.
    """

    default_action: str
    rules: frozenset[CanonicalRule] = field(default_factory=frozenset)
    metadata: dict[str, Any] = field(default_factory=dict)

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, CanonicalPolicy):
            return NotImplemented
        return (
            self.default_action == other.default_action
            and self.rules == other.rules
        )

    def __hash__(self) -> int:
        return hash((self.default_action, self.rules))

    @property
    def has_reconstructed(self) -> bool:
        """True if any rules contain reconstructed predicates."""
        return any(r.reconstructed_predicates for r in self.rules)

    def semantic_eq(self, other: CanonicalPolicy) -> bool:
        """Compare policies using flattened leaf predicate sets.

        This is more permissive than __eq__ - it considers policies equivalent
        if they have the same default action and each rule has the same
        set of leaf predicates, even if grouping structure differs.

        This handles the common case where:
        - Source: (allow op (pred-a) (pred-b))
        - Reversed: (allow op (require-any (pred-a) (pred-b)))
        """
        if self.default_action != other.default_action:
            return False

        self_ops = {r.operation: r for r in self.rules}
        other_ops = {r.operation: r for r in other.rules}

        # All operations must be present in both
        if set(self_ops.keys()) != set(other_ops.keys()):
            return False

        # All rules must be semantically equivalent
        for op in self_ops:
            if not self_ops[op].semantic_eq(other_ops[op]):
                return False

        return True

    def diff(
        self,
        other: CanonicalPolicy,
        *,
        semantic: bool = False,
        tolerate_redundant: bool = False,
        wildcard_equivalence: bool = False,
    ) -> dict[str, Any]:
        """Compute a detailed diff against another CanonicalPolicy.

        Parameters
        ----------
        other : CanonicalPolicy
            The policy to compare against.
        semantic : bool
            If True, use semantic (leaf predicate set) comparison instead of
            structural comparison. This ignores require-any/require-all grouping
            differences.
        tolerate_redundant : bool
            If True, tolerate rules in self that don't appear in other if their
            action matches the default action. This handles cases where source
            has explicit `(deny X)` with a deny default - semantically equivalent
            to not having the rule at all.
        wildcard_equivalence : bool
            If True, recognize wildcard operations as equivalent to their children.
            For example, `user-preference*` in other matches `user-preference-read`
            and `user-preference-write` in self if predicates are compatible.
        """
        if self.default_action != other.default_action:
            return {
                "type": "default_action_mismatch",
                "self": self.default_action,
                "other": other.default_action,
            }

        # Group rules by operation - an operation may have multiple rules
        # (e.g., two appleevent-send rules with different predicates)
        self_ops: dict[str, list[CanonicalRule]] = {}
        for r in self.rules:
            self_ops.setdefault(r.operation, []).append(r)
        other_ops: dict[str, list[CanonicalRule]] = {}
        for r in other.rules:
            other_ops.setdefault(r.operation, []).append(r)

        only_self = set(self_ops.keys()) - set(other_ops.keys())
        only_other = set(other_ops.keys()) - set(self_ops.keys())

        # Filter out redundant rules if requested
        redundant_self: set[str] = set()
        if tolerate_redundant:
            for op in list(only_self):
                rules = self_ops[op]
                # All rules for this op must be redundant (action matches default)
                if all(rule.action == self.default_action for rule in rules):
                    redundant_self.add(op)
            only_self -= redundant_self

        # Handle wildcard equivalence
        wildcard_matches: list[dict[str, Any]] = []
        if wildcard_equivalence:
            for wildcard in list(only_other):
                if wildcard not in WILDCARD_CHILDREN:
                    continue

                children = WILDCARD_CHILDREN[wildcard]
                matched_children = children & only_self

                if not matched_children:
                    continue

                # Get predicates from the wildcard rule(s)
                wildcard_rules = other_ops[wildcard]
                wildcard_preds: set[str] = set()
                for r in wildcard_rules:
                    wildcard_preds.update(str(p) for p in r.predicates)

                # Get predicates from all matched children
                children_preds: set[str] = set()
                for child in matched_children:
                    for r in self_ops[child]:
                        children_preds.update(str(p) for p in r.predicates)

                # Check predicate compatibility:
                # - If wildcard has no predicates (unconditional), it covers everything
                # - If children's predicates are subset of wildcard's, it's compatible
                # - If predicates match exactly, it's compatible
                preds_compatible = (
                    not wildcard_preds or  # Unconditional wildcard
                    children_preds == wildcard_preds or  # Exact match
                    children_preds <= wildcard_preds  # Subset
                )

                if preds_compatible:
                    wildcard_matches.append({
                        "wildcard": wildcard,
                        "matched_children": sorted(matched_children),
                        "wildcard_predicates": sorted(wildcard_preds),
                        "children_predicates": sorted(children_preds),
                    })
                    only_other.discard(wildcard)
                    only_self -= matched_children

        common = set(self_ops.keys()) & set(other_ops.keys())

        predicate_diffs = []
        for op in common:
            self_rules = self_ops[op]
            other_rules = other_ops[op]

            # Compare rule sets for this operation
            if semantic:
                # Semantic comparison: compare flattened leaf predicate sets
                # Merge all leaf predicates from all rules for each side
                self_leaves = set()
                for r in self_rules:
                    self_leaves.update(r.leaf_predicates)
                other_leaves = set()
                for r in other_rules:
                    other_leaves.update(r.leaf_predicates)

                if self_leaves != other_leaves:
                    predicate_diffs.append({
                        "operation": op,
                        "self_predicates": sorted(self_leaves),
                        "other_predicates": sorted(other_leaves),
                        "comparison_mode": "semantic",
                    })
            else:
                # Structural comparison: rule sets must match exactly
                self_rule_set = frozenset(self_rules)
                other_rule_set = frozenset(other_rules)

                if self_rule_set != other_rule_set:
                    # Collect all predicates from all rules for reporting
                    self_preds = set()
                    for r in self_rules:
                        self_preds.update(str(p) for p in r.predicates)
                    other_preds = set()
                    for r in other_rules:
                        other_preds.update(str(p) for p in r.predicates)

                    predicate_diffs.append({
                        "operation": op,
                        "self_predicates": sorted(self_preds),
                        "other_predicates": sorted(other_preds),
                        "self_rule_count": len(self_rules),
                        "other_rule_count": len(other_rules),
                    })

        result = {
            "type": "rule_diff",
            "only_self": sorted(only_self),
            "only_other": sorted(only_other),
            "predicate_diffs": predicate_diffs,
        }
        if tolerate_redundant and redundant_self:
            result["redundant_self"] = sorted(redundant_self)
        if wildcard_equivalence and wildcard_matches:
            result["wildcard_matches"] = wildcard_matches
        return result


@dataclass(frozen=True)
class OpSlotInfo:
    """Normalized information about one op-table slot.

    Used for IR-level comparison where we compare slot assignments
    rather than SBPL text.
    """

    operation: str
    assignment_class: str  # "explicit", "default", "implicit_baseline", etc.
    node_offset: int
    equivalence_group: str  # Normalized class for tolerance comparison

    @staticmethod
    def normalize_class(raw_class: str) -> str:
        """Map raw assignment class to equivalence group.

        Classes in the same equivalence group are considered equivalent
        for roundtrip comparison.
        """
        # Implicit family: these are all "compiler-managed" slots
        implicit_family = {
            "default",
            "implicit_baseline",
            "implicit_nonterm_guard",
            "contextual_implicit",
            "bare_opposite_terminal",
            "unknown",
        }
        if raw_class in implicit_family:
            return "implicit"
        return raw_class  # "explicit" stays "explicit"


@dataclass
class CanonicalPolicyIR:
    """Normalized IR representation of a compiled policy.

    This captures the structural information from the compiled blob
    in a form suitable for equivalence comparison.

    Attributes
    ----------
    default_action : str
        The default action from the profile header.
    op_slots : dict[str, OpSlotInfo]
        Normalized op-table slot information, keyed by operation name.
    graph_signatures : dict[str, str]
        Structural hash of the reachable graph for each explicit operation.
    disconnected_filters : list[dict]
        Filters that are relevant to operations but unreachable from
        their entry points.
    equivalence_level : str
        The level of equivalence this IR supports:
        - "strict": Byte-identical roundtrip possible
        - "semantic": Same runtime behavior, different bytes
        - "degraded": Some information lost, behavior may differ
    """

    default_action: str
    op_slots: dict[str, OpSlotInfo] = field(default_factory=dict)
    graph_signatures: dict[str, str] = field(default_factory=dict)
    disconnected_filters: list[dict] = field(default_factory=list)
    equivalence_level: str = "strict"
    metadata: dict[str, Any] = field(default_factory=dict)

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, CanonicalPolicyIR):
            return NotImplemented
        # Compare using equivalence groups, not raw classes
        self_slots = {
            op: (info.equivalence_group, info.node_offset)
            for op, info in self.op_slots.items()
        }
        other_slots = {
            op: (info.equivalence_group, info.node_offset)
            for op, info in other.op_slots.items()
        }
        return (
            self.default_action == other.default_action
            and self_slots == other_slots
            and self.graph_signatures == other.graph_signatures
        )

    def diff(self, other: CanonicalPolicyIR) -> dict[str, Any]:
        """Compute a detailed diff against another CanonicalPolicyIR."""
        slot_diffs = []
        all_ops = set(self.op_slots.keys()) | set(other.op_slots.keys())

        for op in sorted(all_ops):
            s = self.op_slots.get(op)
            o = other.op_slots.get(op)
            if s is None or o is None:
                slot_diffs.append({
                    "operation": op,
                    "reason": "missing" if s is None else "extra",
                })
            elif s.equivalence_group != o.equivalence_group:
                slot_diffs.append({
                    "operation": op,
                    "reason": "class_mismatch",
                    "self_class": s.assignment_class,
                    "other_class": o.assignment_class,
                    "self_group": s.equivalence_group,
                    "other_group": o.equivalence_group,
                })

        graph_diffs = []
        all_graph_ops = set(self.graph_signatures.keys()) | set(other.graph_signatures.keys())
        for op in sorted(all_graph_ops):
            s_sig = self.graph_signatures.get(op)
            o_sig = other.graph_signatures.get(op)
            if s_sig != o_sig:
                graph_diffs.append({
                    "operation": op,
                    "self_signature": s_sig,
                    "other_signature": o_sig,
                })

        return {
            "default_match": self.default_action == other.default_action,
            "slot_diffs": slot_diffs,
            "graph_diffs": graph_diffs,
            "disconnected_count_self": len(self.disconnected_filters),
            "disconnected_count_other": len(other.disconnected_filters),
        }


# ---------------------------------------------------------------------------
# Predicate flattening utilities
# ---------------------------------------------------------------------------

def _flatten_predicate_to_leaves(predicate_text: str) -> set[str]:
    """Flatten a predicate, extracting leaf predicates from require-any/require-all.

    This recursively unwraps grouping constructs to extract the underlying
    leaf predicates. Useful for set-based comparison that ignores grouping.

    Examples:
        "(literal \"/foo\")" -> {"(literal \"/foo\")"}
        "(require-any (literal \"/a\") (literal \"/b\"))" -> {"(literal \"/a\")", "(literal \"/b\")"}
        "(require-all (subpath \"/x\") (extension \"read\"))" -> {"(subpath \"/x\")", "(extension \"read\")"}

    Parameters
    ----------
    predicate_text : str
        A single predicate string, possibly containing nested require-any/require-all.

    Returns
    -------
    set[str]
        Set of leaf predicate strings (no require-any/require-all wrappers).
    """
    predicate_text = predicate_text.strip()
    if not predicate_text.startswith("("):
        return {predicate_text} if predicate_text else set()

    # Check if this is a grouping construct
    for keyword in ("require-any", "require-all"):
        prefix = f"({keyword} "
        if predicate_text.startswith(prefix):
            # Extract and flatten children
            inner = predicate_text[len(prefix):-1]  # Remove outer parens and keyword
            children = _extract_sexp_children(inner)
            leaves: set[str] = set()
            for child in children:
                leaves.update(_flatten_predicate_to_leaves(child))
            return leaves

    # Not a grouping construct - this is a leaf predicate
    return {predicate_text}


def _extract_sexp_children(text: str) -> list[str]:
    """Extract top-level S-expression children from text.

    Handles nested parens and quoted strings correctly.
    """
    children: list[str] = []
    i = 0
    text = text.strip()

    while i < len(text):
        # Skip whitespace
        while i < len(text) and text[i].isspace():
            i += 1
        if i >= len(text):
            break

        if text[i] == '(':
            # Extract balanced s-expression
            expr = _extract_balanced_sexp(text, i)
            if expr:
                children.append(expr)
                i += len(expr)
            else:
                i += 1
        else:
            # Bare token
            start = i
            while i < len(text) and not text[i].isspace() and text[i] not in '()':
                i += 1
            if i > start:
                children.append(text[start:i])

    return children


def _extract_balanced_sexp(text: str, start: int) -> str | None:
    """Extract a balanced s-expression starting at the given position.

    Handles parentheses inside quoted strings (both regular "..." and regex #"...")
    by tracking string context.
    """
    if start >= len(text) or text[start] != '(':
        return None

    depth = 0
    in_string = False
    i = start

    while i < len(text):
        char = text[i]

        # Handle string start/end
        if char == '"' and not in_string:
            in_string = True
            i += 1
            continue
        elif char == '"' and in_string:
            # Check for escape (look back, but handle multiple escapes)
            num_backslashes = 0
            j = i - 1
            while j >= start and text[j] == '\\':
                num_backslashes += 1
                j -= 1
            # Escaped if odd number of backslashes before quote
            if num_backslashes % 2 == 1:
                i += 1
                continue
            in_string = False
            i += 1
            continue

        # Inside string, just advance
        if in_string:
            i += 1
            continue

        # Outside string, handle parens
        if char == '(':
            depth += 1
        elif char == ')':
            depth -= 1
            if depth == 0:
                return text[start:i + 1]

        i += 1

    # Unbalanced
    return None


# ---------------------------------------------------------------------------
# Deny denormalization utilities
# ---------------------------------------------------------------------------

def has_require_not(predicate_text: str) -> bool:
    """Check if a predicate contains require-not.

    Parameters
    ----------
    predicate_text : str
        A predicate string to check.

    Returns
    -------
    bool
        True if the predicate contains a require-not construct.
    """
    return "(require-not " in predicate_text


def extract_require_not_from_require_all(
    predicate_text: str,
) -> tuple[list[str], list[str]]:
    """Extract require-not predicates from a require-all structure.

    Given a predicate like:
        (require-all (require-not (subpath "/secret")) (subpath "/tmp"))

    Returns:
        negated: ["(subpath \"/secret\")"]
        remaining: ["(subpath \"/tmp\")"]

    Parameters
    ----------
    predicate_text : str
        A predicate string, typically a require-all containing require-not.

    Returns
    -------
    tuple[list[str], list[str]]
        (negated_predicates, remaining_predicates)
        - negated_predicates: predicates that were inside require-not (unwrapped)
        - remaining_predicates: predicates that were not negated
    """
    predicate_text = predicate_text.strip()

    # Must be a require-all to decompose
    if not predicate_text.startswith("(require-all "):
        # Not a require-all - check if it's a bare require-not
        if predicate_text.startswith("(require-not "):
            inner = predicate_text[len("(require-not "):-1].strip()
            return [inner], []
        return [], [predicate_text] if predicate_text else []

    # Extract children from require-all
    inner = predicate_text[len("(require-all "):-1]
    children = _extract_sexp_children(inner)

    negated: list[str] = []
    remaining: list[str] = []

    for child in children:
        child = child.strip()
        if child.startswith("(require-not "):
            # Unwrap the require-not
            inner_pred = child[len("(require-not "):-1].strip()
            negated.append(inner_pred)
        else:
            remaining.append(child)

    return negated, remaining


def rebuild_predicate_without_require_not(
    predicate_text: str,
) -> str | None:
    """Remove require-not children from a require-all predicate.

    Given:
        (require-all (require-not (subpath "/secret")) (subpath "/tmp"))

    Returns:
        (subpath "/tmp")

    If only one non-negated predicate remains, returns it directly.
    If multiple remain, wraps in require-all.
    If none remain, returns None.

    Parameters
    ----------
    predicate_text : str
        A predicate string to transform.

    Returns
    -------
    str | None
        The predicate with require-not removed, or None if nothing remains.
    """
    _, remaining = extract_require_not_from_require_all(predicate_text)

    if not remaining:
        return None
    if len(remaining) == 1:
        return remaining[0]
    return "(require-all " + " ".join(remaining) + ")"