Module refinery.lib.scripts.ps1.deobfuscation.unflatten

"""
Recover original control flow from control-flow-flattened PowerShell scripts.

Control flow flattening replaces sequential and branching code with a dispatcher loop: a while loop
containing a single switch on a state variable, where each case sets the state variable to
determine the next case to execute. This transformer identifies the dispatcher pattern, extracts
the state machine, and recovers the original structure.
"""
from __future__ import annotations

from collections import deque
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Callable, Generator, NamedTuple

from refinery.lib.scripts import Block, Node, Statement, Transformer, tree_root
from refinery.lib.scripts.ps1.analysis.cache import Ps1ModelCache, model_cache
from refinery.lib.scripts.ps1.analysis.effects import Ps1BodyReads, reads_in_bodies
from refinery.lib.scripts.ps1.analysis.naming import Ps1NameRole, named_references, reads_unreadable_name
from refinery.lib.scripts.ps1.analysis.opaque import runs_unreadable_code, writes_nobody_can_attribute
from refinery.lib.scripts.ps1.analysis.values import integer_of, read
from refinery.lib.scripts.ps1.ast import binding_key, get_body, is_builtin_variable, unwrap_parens
from refinery.lib.scripts.ps1.data import COMPARISON_OPS
from refinery.lib.scripts.ps1.deobfuscation.emulator import evaluate_truthy
from refinery.lib.scripts.ps1.deobfuscation.helpers import inside_value_producing_context
from refinery.lib.scripts.ps1.deobfuscation.removal import Ps1RemovalPlan
from refinery.lib.scripts.ps1.model import (
    Expression,
    Ps1AssignmentExpression,
    Ps1BinaryExpression,
    Ps1BreakStatement,
    Ps1ExpressionStatement,
    Ps1IfStatement,
    Ps1IntegerLiteral,
    Ps1ParenExpression,
    Ps1RealLiteral,
    Ps1ScopeModifier,
    Ps1Script,
    Ps1StringLiteral,
    Ps1SwitchStatement,
    Ps1UnaryExpression,
    Ps1Variable,
    Ps1WhileLoop,
)

_MAX_STATES = 500
_MAX_UNROLL_ITERATIONS = 500

if TYPE_CHECKING:
    from refinery.lib.scripts.ps1.analysis.callgraph import Ps1CallGraph
    _VarKey = tuple[str, Ps1ScopeModifier]
    _StateKey = int | float | str


def _is_bool_literal(node: Node) -> bool | None:
    """
    Check if a node is a `$True` or `$False` variable literal. Returns the boolean value,
    or `None` if the node is not a boolean literal.
    """
    if is_builtin_variable(node, frozenset({'true', 'false'})):
        return node.name.lower() == 'true'
    return None


def _integer_state_key(node: Ps1IntegerLiteral) -> int | None:
    """
    The integer PowerShell gives an integer literal, or `None` for a numeral too wide to type. This
    is not `node.value`, which is the numeral's magnitude: a hexadecimal literal is typed by its
    width and carries that width's sign, so `0xFFFFFFFF` is the Int32 -1 rather than four billion,
    and a state machine that matches the magnitude routes itself into a case the host never enters.
    The value the host assigns is what `read` types and `integer_of` reads back.
    """
    return integer_of(read(node))


def _unwrap_constant(node) -> _StateKey | None:
    """
    Extract a constant value (int, float, or string) from an AST node.
    """
    node = unwrap_parens(node) if isinstance(node, Expression) else node
    if isinstance(node, Ps1IntegerLiteral):
        return _integer_state_key(node)
    if isinstance(node, Ps1RealLiteral):
        return node.value
    if isinstance(node, Ps1StringLiteral):
        return node.value
    if isinstance(node, Ps1UnaryExpression) and node.operator == '-':
        inner = unwrap_parens(node.operand) if isinstance(node.operand, Expression) else node.operand
        if isinstance(inner, Ps1IntegerLiteral):
            magnitude = _integer_state_key(inner)
            return None if magnitude is None else -magnitude
        if isinstance(inner, Ps1RealLiteral):
            return -inner.value
    if is_builtin_variable(node, frozenset({'null'})):
        return 0
    return None


@dataclass
class _LinearTransition:
    target: _StateKey

    @property
    def successors(self) -> tuple[_StateKey, ...]:
        return (self.target,)


@dataclass
class _ConditionalTransition:
    condition: Expression
    true_target: _StateKey
    false_target: _StateKey
    true_prefix: list[Statement] = field(default_factory=list)
    false_prefix: list[Statement] = field(default_factory=list)

    @property
    def successors(self) -> tuple[_StateKey, ...]:
        return (self.true_target, self.false_target)


@dataclass
class _ExitTransition:

    @property
    def successors(self) -> tuple[_StateKey, ...]:
        return ()


if TYPE_CHECKING:
    _Transition = _LinearTransition | _ConditionalTransition | _ExitTransition


@dataclass
class _StateBlock:
    state_id: _StateKey
    statements: list[Statement]
    transition: _Transition


@dataclass
class _DispatcherMatch:
    state_var_name: str
    state_var_scope: Ps1ScopeModifier
    condition: Expression
    switch: Ps1SwitchStatement


def _var_key(node: Ps1Variable) -> _VarKey:
    return (node.name.lower(), node.scope)


def _collect_variables(node: Node) -> set[_VarKey]:
    """
    Collect all distinct variable keys referenced in an expression tree.
    """
    result: set[_VarKey] = set()
    for child in node.walk():
        if isinstance(child, Ps1Variable):
            if child.scope != Ps1ScopeModifier.NONE:
                continue
            lower = child.name.lower()
            if lower in ('true', 'false', 'null'):
                continue
            result.add(_var_key(child))
    return result


def _make_exit_check(
    condition: Expression,
    var_name: str,
) -> Callable[[_StateKey], bool]:
    """
    Return a predicate that checks whether assigning a given state value to the state variable
    would make the while condition falsy (the loop would exit). Uses the emulator to evaluate the
    condition with the state variable bound.

    A condition the emulator cannot evaluate is read as one that does not exit, which is sound only
    because `_exit_condition_is_total` refuses the whole recovery unless the condition resolves at
    every state value the predicate is consulted on. Without that gate a loop whose exit depends on
    an unevaluable value — `while ($s -ne (Get-Random))` — would recover as one that never exits.
    """
    def is_exit(state_value: _StateKey) -> bool:
        bindings = {var_name: state_value}
        result = evaluate_truthy(condition, bindings)
        if result is None:
            return False
        return not result
    return is_exit


def _exit_condition_is_total(
    condition: Expression,
    var_name: str,
    states: dict[_StateKey, _StateBlock],
) -> bool:
    """
    Whether the loop condition resolves to a definite truth value at every state the machine can
    hold — each case label and each transition target, which together are exactly the values
    `_make_exit_check`'s predicate is asked about. A machine undecidable at one of them cannot be
    shown to exit there, so recovering it would run a body the host may never run.
    """
    values: set[_StateKey] = set(states)
    for block in states.values():
        values.update(block.transition.successors)
    return all(
        evaluate_truthy(condition, {var_name: value}) is not None
        for value in values
    )


def _is_state_assignment(
    stmt: Statement,
    var_name: str,
    var_scope: Ps1ScopeModifier,
) -> _StateKey | None:
    """
    If the statement is an assignment to the state variable with a constant value, return that
    value. Otherwise return `None`.
    """
    result = _get_simple_assignment(stmt)
    if result is None:
        return None
    key, value = result
    if key != (var_name, var_scope):
        return None
    return _unwrap_constant(value)


def _get_simple_assignment(
    stmt: Statement,
) -> tuple[_VarKey, Node] | None:
    """
    If the statement is a simple assignment ($var = <value>), return the variable key
    and the value node. Returns `None` for non-assignments or compound operators.
    """
    if not isinstance(stmt, Ps1ExpressionStatement):
        return None
    expr = stmt.expression
    if not isinstance(expr, Ps1AssignmentExpression):
        return None
    if expr.operator != '=':
        return None
    target = expr.target
    if not isinstance(target, Ps1Variable):
        return None
    if expr.value is None:
        return None
    return (_var_key(target), expr.value)


def _resolve_value(
    node: Node,
    env: dict[_VarKey, _StateKey | bool],
) -> _StateKey | None:
    if not isinstance(node, Expression):
        return None
    expr = unwrap_parens(node)
    value = _unwrap_constant(expr)
    if value is not None:
        return value
    if isinstance(expr, Ps1Variable):
        key = _var_key(expr)
        val = env.get(key)
        if isinstance(val, (int, float, str)):
            return val
    return None


def _resolve_bool(
    node: Node,
    env: dict[_VarKey, _StateKey | bool],
) -> bool | None:
    if not isinstance(node, Expression):
        return None
    expr = unwrap_parens(node)
    literal = _is_bool_literal(expr)
    if literal is not None:
        return literal
    if isinstance(expr, Ps1Variable):
        key = _var_key(expr)
        val = env.get(key)
        if isinstance(val, bool):
            return val
    if (
        isinstance(expr, Ps1BinaryExpression)
        and (_l := expr.left) is not None
        and (_r := expr.right) is not None
    ):
        if comparator := COMPARISON_OPS.get(op := expr.operator.lower()):
            lhs = _resolve_value(_l, env)
            rhs = _resolve_value(_r, env)
            if lhs is not None and rhs is not None:
                if type(lhs) is not type(rhs):
                    return None
                return comparator(lhs, rhs)
        else:
            lhs = _resolve_bool(_l, env)
            rhs = _resolve_bool(_r, env)
            if lhs is None:
                lhs, rhs = rhs, lhs
            if op == '-and':
                return lhs and rhs
            if op == '-or':
                return lhs or rhs
    if (
        isinstance(expr, Ps1UnaryExpression)
        and (_v := expr.operand) is not None
        and (op := expr.operator.lower()) in ('-not', '!')
        and (bv := _resolve_bool(_v, env)) is not None
    ):
        return not bv
    return None


def _match_dispatcher(loop: Ps1WhileLoop) -> _DispatcherMatch | None:
    """
    Check whether a while loop matches the CFF dispatcher pattern: a while loop whose condition
    involves a single variable, and whose body is a single switch on that variable.
    """
    cond = loop.condition
    if cond is None:
        return None
    body = loop.body
    if body is None or len(body.body) != 1:
        return None
    switch = body.body[0]
    if not isinstance(switch, Ps1SwitchStatement):
        return None
    if switch.file:
        return None
    switch_val = switch.value
    if not isinstance(switch_val, Ps1Variable):
        return None
    var_name = switch_val.name.lower()
    var_scope = switch_val.scope
    cond_vars = _collect_variables(cond)
    if not cond_vars:
        return None
    state_key = (var_name, var_scope)
    if state_key not in cond_vars:
        return None
    if len(cond_vars) > 1:
        return None
    return _DispatcherMatch(
        state_var_name=var_name,
        state_var_scope=var_scope,
        condition=cond,
        switch=switch,
    )


def _find_state_init(
    body: list[Statement],
    loop_index: int,
    var_name: str,
    var_scope: Ps1ScopeModifier,
    calls: Ps1CallGraph,
) -> tuple[int, _StateKey] | None:
    """
    Scan backwards from the while loop to the statement that seeds the state variable with a
    constant, and return its index and the constant; `None` where there is none, or where a
    statement in between may write the variable (`_may_write_the_state`), since the loop then does
    not start from the seed.
    """
    for i in range(loop_index - 1, -1, -1):
        value = _is_state_assignment(body[i], var_name, var_scope)
        if value is not None:
            return (i, value)
        if _may_write_the_state(body[i], var_name, calls):
            return None
    return None


def _may_write_the_state(statement: Node, var_name: str, calls: Ps1CallGraph) -> bool:
    """
    Whether running *statement* may store into the state variable *var_name*: a variable of that
    name stands in it under any scope, a command in it writes the name it spells
    (`refinery.lib.scripts.ps1.analysis.naming.named_references`), something in it writes where
    no reading of the source names the target
    (`refinery.lib.scripts.ps1.analysis.opaque.writes_nobody_can_attribute`), or something in it
    may run a body of the script (`Ps1CallGraph.may_run_a_body`), which may write any name.
    """
    for node in statement.walk():
        if isinstance(node, Ps1Variable) and binding_key(node) == var_name:
            return True
        if writes_nobody_can_attribute(node) or calls.may_run_a_body(node):
            return True
        if any(
            reference.key == var_name and reference.role is not Ps1NameRole.READS
            for reference in named_references(node)
        ):
            return True
    return False


class _StateReads:
    """
    Which statements of one body may read which variables, read once for the body: a statement may
    read a variable of a name that stands in it under any scope or that a command in it names;
    where something in it may run a body of the script (`Ps1CallGraph.may_run_a_body`), a variable
    a function or a stored script block reads (*bodies*); and every variable where it reads a name
    no reading of the source spells or runs code nobody reads.
    """

    def __init__(self, body: list[Statement], calls: Ps1CallGraph, bodies: Ps1BodyReads):
        self._first: dict[str, int] = {}
        self._last: dict[str, int] = {}
        self._every: list[int] = []
        self._running: list[int] = []
        self._read_in_bodies = bodies.names
        for index, statement in enumerate(body):
            every = running = False
            for node in statement.walk():
                if isinstance(node, Ps1Variable):
                    self._mark(binding_key(node), index)
                for reference in named_references(node):
                    self._mark(reference.key, index)
                if reads_unreadable_name(node) or runs_unreadable_code(node):
                    every = True
                elif calls.may_run_a_body(node):
                    running = True
            if every or running and bodies.dynamic:
                self._every.append(index)
            elif running:
                self._running.append(index)

    def _mark(self, name: str, index: int) -> None:
        self._first.setdefault(name, index)
        self._last[name] = index

    def _runs_a_reader(self, var_name: str) -> list[int]:
        return self._running if var_name in self._read_in_bodies else []

    def after(self, index: int, var_name: str) -> bool:
        """
        Whether a statement after the one at *index* may read *var_name*.
        """
        later = [self._last.get(var_name, -1), *self._every[-1:], *self._runs_a_reader(var_name)[-1:]]
        return max(later) > index

    def first(self, var_name: str) -> int | None:
        """
        The index of the first statement that may read *var_name*, or `None` where none may.
        """
        found = [*self._every[:1], *self._runs_a_reader(var_name)[:1]]
        if (mentioned := self._first.get(var_name)) is not None:
            found.append(mentioned)
        return min(found, default=None)


def _strip_trailing_break(stmts: list[Statement]) -> list[Statement]:
    """
    Remove a trailing break statement (part of the switch dispatch, not the original code).
    """
    if stmts and isinstance(stmts[-1], Ps1BreakStatement) and stmts[-1].label is None:
        return stmts[:-1]
    return stmts


def _extract_transition(
    stmts: list[Statement],
    var_name: str,
    var_scope: Ps1ScopeModifier,
    is_exit: Callable[[_StateKey], bool],
) -> tuple[list[Statement], _Transition] | None:
    """
    Separate a switch case body into side-effect statements and a state transition. Returns
    (side_effects, transition) or `None` if the pattern is not recognized.
    """
    if not stmts:
        return None
    last = stmts[-1]
    state_val = _is_state_assignment(last, var_name, var_scope)
    if state_val is not None:
        side_effects = list(stmts[:-1])
        if is_exit(state_val):
            return (side_effects, _ExitTransition())
        return (side_effects, _LinearTransition(target=state_val))
    if isinstance(last, Ps1IfStatement) and last.else_block is not None and len(last.clauses) == 1:
        condition, true_block = last.clauses[0]
        false_block = last.else_block
        true_body = _strip_trailing_break(list(true_block.body))
        false_body = _strip_trailing_break(list(false_block.body))
        if not true_body or not false_body:
            return None
        true_state = _is_state_assignment(true_body[-1], var_name, var_scope)
        false_state = _is_state_assignment(false_body[-1], var_name, var_scope)
        if true_state is None or false_state is None:
            return None
        true_prefix = true_body[:-1]
        false_prefix = false_body[:-1]
        side_effects = list(stmts[:-1])
        return (side_effects, _ConditionalTransition(
            condition=condition,
            true_target=true_state,
            false_target=false_state,
            true_prefix=true_prefix,
            false_prefix=false_prefix,
        ))
    return None


def _extract_state_machine(
    match: _DispatcherMatch,
    is_exit: Callable[[_StateKey], bool],
) -> dict[_StateKey, _StateBlock] | None:
    """
    Parse all switch cases into a state machine dictionary. Returns `None` on failure.
    """
    states: dict[_StateKey, _StateBlock] = {}
    var_name = match.state_var_name
    var_scope = match.state_var_scope
    for condition, block in match.switch.clauses:
        if condition is None:
            continue
        state_id = _unwrap_constant(condition)
        if state_id is None:
            return None
        body = _strip_trailing_break(list(block.body))
        result = _extract_transition(body, var_name, var_scope, is_exit)
        if result is None:
            return None
        side_effects, transition = result
        states[state_id] = _StateBlock(
            state_id=state_id,
            statements=side_effects,
            transition=transition,
        )
    if len(states) > _MAX_STATES:
        return None
    return states


def _negate_condition(cond: Expression) -> Expression:
    """
    Return the logical negation of a condition expression. Tries to simplify where possible (e.g.,
    flipping -eq to -ne) rather than wrapping in -Not.
    """
    unwrapped = unwrap_parens(cond)
    if not isinstance(unwrapped, Expression):
        return Ps1UnaryExpression(operator='-Not', operand=cond, prefix=True)
    cond = unwrapped
    if isinstance(cond, Ps1BinaryExpression):
        flipped = {
            '-eq': '-NE',
            '-ne': '-EQ',
            '-lt': '-GE',
            '-ge': '-LT',
            '-gt': '-LE',
            '-le': '-GT',
        }.get(cond.operator.lower())
        if flipped is not None:
            return Ps1BinaryExpression(
                left=cond.left,
                operator=flipped,
                right=cond.right,
            )
    if (
        isinstance(cond, Ps1UnaryExpression)
        and cond.operator.lower() in ('-not', '!')
        and cond.operand is not None
    ):
        return cond.operand
    return Ps1UnaryExpression(operator='-Not', operand=Ps1ParenExpression(expression=cond))


def _build_if(
    condition: Expression,
    true_body: list[Statement],
    false_body: list[Statement],
) -> Ps1IfStatement:
    """
    Build an if/else statement. Omits the else block if the false body is empty; negates the
    condition if only the true body is empty. Where both are empty, the statement still evaluates
    its condition, which may write or fail, so it is kept with an empty body.
    """
    if not true_body and not false_body:
        return Ps1IfStatement(clauses=[(condition, Block(body=[]))])
    if not true_body:
        return Ps1IfStatement(
            clauses=[(_negate_condition(condition), Block(body=false_body))],
        )
    if not false_body:
        return Ps1IfStatement(
            clauses=[(condition, Block(body=true_body))],
        )
    return Ps1IfStatement(
        clauses=[(condition, Block(body=true_body))],
        else_block=Block(body=false_body),
    )


def _find_back_edges(
    states: dict[_StateKey, _StateBlock],
    entry: _StateKey,
    is_exit: Callable[[_StateKey], bool],
) -> dict[_StateKey, _StateKey]:
    """
    Walk the state graph from the entry and find back-edges using DFS. Returns a mapping from back-
    edge source state to back-edge target state (the loop header).
    """
    back_edges: dict[_StateKey, _StateKey] = {}
    visiting: set[_StateKey] = set()
    visited: set[_StateKey] = set()

    def _dfs(state: _StateKey):
        if is_exit(state) or state not in states:
            return
        if state in visited:
            return
        visiting.add(state)
        block = states[state]
        for target in block.transition.successors:
            if is_exit(target) or target not in states:
                continue
            if target in visiting:
                back_edges[state] = target
            elif target not in visited:
                _dfs(target)
        visiting.discard(state)
        visited.add(state)

    _dfs(entry)
    return back_edges


def _find_join_point(
    states: dict[_StateKey, _StateBlock],
    true_start: _StateKey,
    false_start: _StateKey,
    is_exit: Callable[[_StateKey], bool],
    back_edge_targets: set[_StateKey],
) -> _StateKey | None:
    """
    Find the first state reachable from both arms of a conditional (the join point). Returns `None`
    if no common state is found (one or both arms exit).
    """
    true_reach: list[_StateKey] = []
    false_reach: list[_StateKey] = []
    true_queue: deque[_StateKey] = deque([true_start])
    false_queue: deque[_StateKey] = deque([false_start])
    true_seen: set[_StateKey] = set()
    false_seen: set[_StateKey] = set()

    def _expand(
        queue: deque[_StateKey],
        seen: set[_StateKey],
        reach: list[_StateKey],
    ):
        while queue:
            s = queue.popleft()
            if is_exit(s) or s not in states:
                continue
            if s in seen:
                continue
            if s in back_edge_targets:
                reach.append(s)
                continue
            seen.add(s)
            reach.append(s)
            for succ in states[s].transition.successors:
                queue.append(succ)

    _expand(true_queue, true_seen, true_reach)
    _expand(false_queue, false_seen, false_reach)
    true_set = set(true_reach)
    for s in false_reach:
        if s in true_set:
            return s
    return None


def _collect_loop_states(
    states: dict[_StateKey, _StateBlock],
    header: _StateKey,
    is_exit: Callable[[_StateKey], bool],
    latches: set[_StateKey],
) -> set[_StateKey]:
    """
    Collect all state IDs that belong to a loop (reachable from header without leaving through
    latches).
    """
    loop_states: set[_StateKey] = set()
    queue: deque[_StateKey] = deque([header])
    while queue:
        s = queue.popleft()
        if is_exit(s) or s not in states or s in loop_states:
            continue
        loop_states.add(s)
        if s in latches and s != header:
            continue
        queue.extend(states[s].transition.successors)
    return loop_states


def _collect_internal_vars(
    states: dict[_StateKey, _StateBlock],
    loop_states: set[_StateKey],
    state_var_key: _VarKey,
) -> set[_VarKey]:
    """
    Identify internal dispatch variables within the loop. These are variables that are only
    assigned constant values (integers, floats, strings, or booleans) or copies of other internal
    variables. They are artifacts of the flattening and should be suppressed in output.
    """
    candidates: set[_VarKey] = {state_var_key}
    changed = True
    while changed:
        changed = False
        for sid in loop_states:
            block = states[sid]
            for stmts in _all_statement_lists(block):
                for stmt in stmts:
                    result = _get_simple_assignment(stmt)
                    if result is None:
                        continue
                    key, value = result
                    if key in candidates:
                        continue
                    if isinstance(value, Expression):
                        value = unwrap_parens(value)
                    if _unwrap_constant(value) is not None:
                        candidates.add(key)
                        changed = True
                    elif isinstance(value, Ps1Variable):
                        if _is_bool_literal(value) is not None:
                            candidates.add(key)
                            changed = True
                        elif _var_key(value) in candidates:
                            candidates.add(key)
                            changed = True
    # A variable referenced by an emitted (non-suppressed) statement is real data, not a dispatch
    # artifact: suppressing its assignment would leave the recovered code reading an undefined
    # variable. Drop such variables from the internal set (the state variable always stays internal).
    used_by_emitted: set[_VarKey] = set()
    for sid in loop_states:
        block = states[sid]
        for stmts in _all_statement_lists(block):
            for stmt in stmts:
                result = _get_simple_assignment(stmt)
                if result is not None and result[0] in candidates:
                    continue
                for node in stmt.walk():
                    if isinstance(node, Ps1Variable):
                        used_by_emitted.add(_var_key(node))
    candidates -= used_by_emitted
    candidates.add(state_var_key)
    return candidates


def _all_statement_lists(block: _StateBlock) -> Generator[list[Statement], None, None]:
    """
    Yield all statement lists within a state block (main statements plus conditional transition
    prefixes).
    """
    yield block.statements
    t = block.transition
    if isinstance(t, _ConditionalTransition):
        yield t.true_prefix
        yield t.false_prefix


def _update_env(
    env: dict[_VarKey, _StateKey | bool],
    key: _VarKey,
    value: Node,
):
    if isinstance(value, Expression):
        value = unwrap_parens(value)
    const_val = _unwrap_constant(value)
    if const_val is not None:
        env[key] = const_val
        return
    if isinstance(value, Ps1Variable):
        bool_val = _is_bool_literal(value)
        if bool_val is not None:
            env[key] = bool_val
            return
        src_key = _var_key(value)
        if src_key in env:
            env[key] = env[src_key]
            return
    env.pop(key, None)


def _simulate_statements(
    stmts: list[Statement],
    env: dict[_VarKey, _StateKey | bool],
    internal_vars: set[_VarKey],
    output: list[Statement],
):
    """
    Process a list of statements during simulation. Updates env for internal assignments, and
    appends non-internal statements to output.
    """
    for stmt in stmts:
        result = _get_simple_assignment(stmt)
        if result is not None:
            key, value = result
            if key in internal_vars:
                _update_env(env, key, value)
                continue
        output.append(stmt)


def _simulate_arm(
    states: dict[_StateKey, _StateBlock],
    start: _StateKey,
    header: _StateKey,
    is_exit: Callable[[_StateKey], bool],
    latches: set[_StateKey],
    env: dict[_VarKey, _StateKey | bool],
    internal_vars: set[_VarKey],
    output: list[Statement],
) -> tuple[list[Statement], _StateKey | None] | None:
    """
    Simulate execution of a single arm (branch) during loop unrolling. Follows the state graph
    linearly, resolving conditions where possible. Returns (statements, next_state) or `None` if a
    condition cannot be resolved.
    """
    current: _StateKey | None = start
    step_count = 0
    while current is not None and step_count < _MAX_STATES * 2:
        step_count += 1
        if is_exit(current):
            return (output, None)
        if current not in states:
            return None
        if current == header:
            return (output, header)
        if current in latches:
            cur_block = states[current]
            _simulate_statements(cur_block.statements, env, internal_vars, output)
            t = cur_block.transition
            if isinstance(t, _LinearTransition) and t.target == header:
                return (output, header)
            return None

        cur_block = states[current]
        _simulate_statements(cur_block.statements, env, internal_vars, output)
        t = cur_block.transition

        if isinstance(t, _ExitTransition):
            return (output, None)
        if isinstance(t, _LinearTransition):
            current = t.target
            continue
        if isinstance(t, _ConditionalTransition):
            branch = _resolve_bool(t.condition, env)
            if branch is True:
                _simulate_statements(t.true_prefix, env, internal_vars, output)
                current = t.true_target
            elif branch is False:
                _simulate_statements(t.false_prefix, env, internal_vars, output)
                current = t.false_target
            else:
                return None
            continue
        return None
    return None


def _seed_env_from_preamble(
    states: dict[_StateKey, _StateBlock],
    entry: _StateKey,
    header: _StateKey,
    is_exit: Callable[[_StateKey], bool],
    internal_vars: set[_VarKey],
) -> dict[_VarKey, _StateKey | bool] | None:
    """
    Walk the linear chain from entry to header, simulating assignments to build the initial
    variable environment. Returns the env or `None` if the path from entry to header is not a
    simple linear chain.
    """
    env: dict[_VarKey, _StateKey | bool] = {}
    discard: list[Statement] = []
    current = entry
    visited: set[_StateKey] = set()
    while current != header:
        if is_exit(current) or current not in states or current in visited:
            return None
        visited.add(current)
        block = states[current]
        _simulate_statements(block.statements, env, internal_vars, discard)
        if isinstance(block.transition, _LinearTransition):
            current = block.transition.target
        else:
            return None
    return env


def _merge_next_state(
    true_next: _StateKey | None,
    false_next: _StateKey | None,
) -> _StateKey | None:
    if true_next == false_next:
        return true_next
    if true_next is None:
        return false_next
    if false_next is None:
        return true_next
    return true_next


def _try_unroll_loop(
    states: dict[_StateKey, _StateBlock],
    entry: _StateKey,
    header: _StateKey,
    is_exit: Callable[[_StateKey], bool],
    back_edges: dict[_StateKey, _StateKey],
    loop_headers: set[_StateKey],
) -> tuple[list[Statement], set[_VarKey]] | None:
    """
    Attempt to unroll a loop by symbolically executing it. Returns

        (unrolled_statements, internal_vars)

    if successful, or `None` if the loop cannot be fully resolved.
    """
    block = states[header]
    if not isinstance(block.transition, _ConditionalTransition):
        return None
    cond_trans = block.transition

    latches: set[_StateKey] = set()
    for latch, target in back_edges.items():
        if target == header:
            latches.add(latch)

    loop_states = _collect_loop_states(states, header, is_exit, latches)

    cond = cond_trans.condition
    loop_var_key: _VarKey | None = None
    cond_unwrapped = unwrap_parens(cond)
    if isinstance(cond_unwrapped, Ps1BinaryExpression):
        left = unwrap_parens(cond_unwrapped.left) if cond_unwrapped.left is not None else None
        right = (
            unwrap_parens(cond_unwrapped.right) if cond_unwrapped.right is not None else None
        )
        if isinstance(left, Ps1Variable):
            loop_var_key = _var_key(left)
        elif isinstance(right, Ps1Variable):
            loop_var_key = _var_key(right)

    if loop_var_key is None:
        return None

    internal_vars = _collect_internal_vars(states, loop_states, loop_var_key)

    direction = _determine_loop_direction(states, cond_trans, header, is_exit, loop_headers)
    if direction is None:
        return None
    body_entry = direction.body_entry
    body_prefix_stmts = direction.prefix

    seed = _seed_env_from_preamble(states, entry, header, is_exit, internal_vars)
    env: dict[_VarKey, _StateKey | bool] = seed if seed is not None else {}
    result: list[Statement] = []

    for iteration in range(_MAX_UNROLL_ITERATIONS):
        cond_result = _resolve_bool(cond, env)
        if cond_result is False:
            return (result, internal_vars)
        if cond_result is None and iteration > 0:
            return None

        iteration_stmts: list[Statement] = []
        _simulate_statements(block.statements, env, internal_vars, iteration_stmts)
        _simulate_statements(body_prefix_stmts, env, internal_vars, iteration_stmts)

        current: _StateKey | None = body_entry
        step_count = 0
        while current is not None and step_count < _MAX_STATES * 2:
            step_count += 1
            if is_exit(current):
                result.extend(iteration_stmts)
                return (result, internal_vars)
            if current not in states:
                return None
            if current == header:
                break
            if current in latches:
                cur_block = states[current]
                _simulate_statements(
                    cur_block.statements, env, internal_vars, iteration_stmts,
                )
                t = cur_block.transition
                if isinstance(t, _LinearTransition) and t.target == header:
                    break
                return None

            cur_block = states[current]
            _simulate_statements(cur_block.statements, env, internal_vars, iteration_stmts)
            t = cur_block.transition

            if isinstance(t, _ExitTransition):
                result.extend(iteration_stmts)
                return (result, internal_vars)

            if isinstance(t, _LinearTransition):
                current = t.target
                continue

            if isinstance(t, _ConditionalTransition):
                branch = _resolve_bool(t.condition, env)
                if branch is True:
                    _simulate_statements(
                        t.true_prefix, env, internal_vars, iteration_stmts,
                    )
                    current = t.true_target
                    continue
                if branch is False:
                    _simulate_statements(
                        t.false_prefix, env, internal_vars, iteration_stmts,
                    )
                    current = t.false_target
                    continue
                true_stmts: list[Statement] = []
                false_stmts: list[Statement] = []
                true_env = dict(env)
                false_env = dict(env)
                _simulate_statements(
                    t.true_prefix, true_env, internal_vars, true_stmts,
                )
                _simulate_statements(
                    t.false_prefix, false_env, internal_vars, false_stmts,
                )
                true_arm_result = _simulate_arm(
                    states,
                    t.true_target,
                    header,
                    is_exit,
                    latches,
                    true_env,
                    internal_vars,
                    true_stmts,
                )
                false_arm_result = _simulate_arm(
                    states,
                    t.false_target,
                    header,
                    is_exit,
                    latches,
                    false_env,
                    internal_vars,
                    false_stmts,
                )
                if true_arm_result is None or false_arm_result is None:
                    return None
                true_arm_body, true_next = true_arm_result
                false_arm_body, false_next = false_arm_result
                iteration_stmts.append(_build_if(t.condition, true_arm_body, false_arm_body))
                for key in set(env.keys()) | set(true_env.keys()) | set(false_env.keys()):
                    tv = true_env.get(key)
                    fv = false_env.get(key)
                    if tv == fv and tv is not None:
                        env[key] = tv
                    else:
                        env.pop(key, None)
                current = _merge_next_state(true_next, false_next)
                continue
            return None

        result.extend(iteration_stmts)

    return None


def _recover_structure(
    states: dict[_StateKey, _StateBlock],
    entry: _StateKey,
    is_exit: Callable[[_StateKey], bool],
) -> list[Statement] | None:
    """
    Walk the state graph from the entry state and emit the recovered AST. Returns `None` if the
    structure cannot be recovered (e.g., irreducible control flow).
    """
    back_edges = _find_back_edges(states, entry, is_exit)
    loop_headers: set[_StateKey] = set(back_edges.values())

    def _emit_arm(
        start: _StateKey,
        stop: _StateKey | None,
        claimed: set[_StateKey],
    ) -> tuple[list[Statement], _StateKey | None] | None:
        """
        Emit statements from start, stopping before the stop state (if given). Returns

            (statements, next_state)

        where next_state is the stop state or `None`.
        """
        result: list[Statement] = []
        current: _StateKey | None = start

        while current is not None:
            if is_exit(current):
                return (result, None)
            if current not in states:
                return None
            if stop is not None and current == stop:
                return (result, current)
            if current in claimed:
                return (result, current)
            if current in loop_headers:
                loop_result = _emit_loop(current, claimed)
                if loop_result is None:
                    return None
                loop_stmts, loop_next, loop_internals = loop_result
                if loop_internals is not None:
                    filtered: list[Statement] = []
                    for s in result:
                        assignment = _get_simple_assignment(s)
                        if assignment is None or assignment[0] not in loop_internals:
                            filtered.append(s)
                    result = filtered
                result.extend(loop_stmts)
                current = loop_next
                continue
            claimed.add(current)
            block = states[current]
            result.extend(block.statements)
            transition = block.transition

            if isinstance(transition, _ExitTransition):
                return (result, None)

            if isinstance(transition, _LinearTransition):
                current = transition.target
                continue

            if isinstance(transition, _ConditionalTransition):
                true_target = transition.true_target
                false_target = transition.false_target
                join = _find_join_point(
                    states, true_target, false_target, is_exit, loop_headers,
                )
                effective_stop = join if join is not None else stop
                true_result = _emit_arm(true_target, effective_stop, claimed)
                false_result = _emit_arm(false_target, effective_stop, claimed)
                if true_result is None or false_result is None:
                    return None
                true_stmts, true_next = true_result
                false_stmts, false_next = false_result
                true_stmts = list(transition.true_prefix) + true_stmts
                false_stmts = list(transition.false_prefix) + false_stmts
                result.append(_build_if(transition.condition, true_stmts, false_stmts))
                current = _merge_next_state(true_next, false_next)
                continue
        return (result, None)

    def _emit_loop(
        header: _StateKey,
        outer_claimed: set[_StateKey],
    ) -> tuple[list[Statement], _StateKey | None, set[_VarKey] | None] | None:
        """
        Emit a while loop rooted at the given header state. First attempts to fully unroll the loop
        via symbolic execution. Falls back to structural recovery if unrolling fails. Returns:

            (statements, next_state, internal_vars_or_None)
        """
        unrolled = _try_unroll_loop(
            states, entry, header, is_exit, back_edges, loop_headers,
        )
        if unrolled is not None:
            unrolled_stmts, unrolled_internals = unrolled
            for sid in _collect_loop_states(
                states,
                header,
                is_exit,
                {latch for latch, target in back_edges.items() if target == header},
            ):
                outer_claimed.add(sid)
            return (unrolled_stmts, None, unrolled_internals)

        block = states[header]
        loop_cond: Expression
        body_start: _StateKey
        exit_target: _StateKey | None = None

        if isinstance(block.transition, _ConditionalTransition):
            cond_trans = block.transition
            direction = _determine_loop_direction(
                states, cond_trans, header, is_exit, loop_headers,
            )
            if direction is not None:
                loop_cond = (
                    _negate_condition(cond_trans.condition) if direction.negated
                    else cond_trans.condition
                )
                body_start_stmts = list(block.statements) + direction.prefix
                body_start = direction.body_entry
                exit_target = direction.exit_target
            else:
                loop_cond = Ps1Variable(name='True', scope=Ps1ScopeModifier.NONE)
                body_start_stmts = list(block.statements)
                body_start = header
                outer_claimed.add(header)
                inner_claimed: set[_StateKey] = set()
                inner_claimed.add(header)
                body_result = _emit_arm(cond_trans.true_target, header, inner_claimed)
                if body_result is None:
                    return None
                body_stmts = body_start_stmts
                if_true, _ = body_result
                if_false_result = _emit_arm(cond_trans.false_target, header, inner_claimed)
                if if_false_result is None:
                    return None
                if_false, _ = if_false_result
                true_body = list(cond_trans.true_prefix) + if_true
                false_body = list(cond_trans.false_prefix) + if_false
                body_stmts.append(_build_if(cond_trans.condition, true_body, false_body))
                while_stmt = Ps1WhileLoop(
                    condition=loop_cond,
                    body=Block(body=body_stmts),
                )
                outer_claimed.update(inner_claimed)
                return ([while_stmt], exit_target, None)
        else:
            loop_cond = Ps1Variable(name='True', scope=Ps1ScopeModifier.NONE)
            body_start_stmts = list(block.statements)
            if isinstance(block.transition, _LinearTransition):
                body_start = block.transition.target
            else:
                return None
            exit_target = None

        outer_claimed.add(header)
        inner_claimed = set()
        inner_claimed.add(header)
        body_result = _emit_arm(body_start, header, inner_claimed)
        if body_result is None:
            return None
        body_stmts, _ = body_result
        all_body = body_start_stmts + body_stmts
        while_stmt = Ps1WhileLoop(
            condition=loop_cond,
            body=Block(body=all_body),
        )
        outer_claimed.update(inner_claimed)
        next_state: _StateKey | None = None
        if exit_target is not None and not is_exit(exit_target):
            next_state = exit_target
        return ([while_stmt], next_state, None)

    claimed: set[_StateKey] = set()
    result = _emit_arm(entry, None, claimed)
    if result is None:
        return None
    stmts, _ = result
    return stmts


class _LoopDirection(NamedTuple):
    body_entry: _StateKey
    exit_target: _StateKey | None
    prefix: list[Statement]
    negated: bool


def _determine_loop_direction(
    states: dict[_StateKey, _StateBlock],
    cond_trans: _ConditionalTransition,
    header: _StateKey,
    is_exit: Callable[[_StateKey], bool],
    loop_headers: set[_StateKey],
) -> _LoopDirection | None:
    """
    Determine which branch of a conditional loop header leads to the loop body
    and which leads to the exit. Returns `None` when direction is ambiguous (both
    branches reach the header, or neither does).
    """
    true_target = cond_trans.true_target
    false_target = cond_trans.false_target
    true_is_exit = is_exit(true_target) or true_target not in states
    false_is_exit = is_exit(false_target) or false_target not in states
    true_reaches = _state_reaches(
        states, true_target, header, is_exit, loop_headers - {header},
    )
    false_reaches = _state_reaches(
        states, false_target, header, is_exit, loop_headers - {header},
    )
    if true_reaches and (not false_reaches or false_is_exit):
        return _LoopDirection(true_target, false_target, list(cond_trans.true_prefix), False)
    if false_reaches and (not true_reaches or true_is_exit):
        return _LoopDirection(false_target, true_target, list(cond_trans.false_prefix), True)
    if false_is_exit:
        return _LoopDirection(true_target, false_target, list(cond_trans.true_prefix), False)
    if true_is_exit:
        return _LoopDirection(false_target, true_target, list(cond_trans.false_prefix), True)
    return None


def _state_reaches(
    states: dict[_StateKey, _StateBlock],
    start: _StateKey,
    target: _StateKey,
    is_exit: Callable[[_StateKey], bool],
    barriers: set[_StateKey],
) -> bool:
    """
    Check if the state graph can reach target from start without crossing barriers.
    """
    visited: set[_StateKey] = set()
    queue: deque[_StateKey] = deque([start])
    while queue:
        s = queue.popleft()
        if s == target:
            return True
        if is_exit(s) or s not in states or s in visited or s in barriers:
            continue
        visited.add(s)
        queue.extend(states[s].transition.successors)
    return False


class Ps1ControlFlowDeflattening(Transformer):
    """
    Recover original control flow from control-flow-flattened scripts.
    """

    def visit(self, node: Node):
        for parent in list(node.walk()):
            if inside_value_producing_context(parent):
                continue
            body = get_body(parent)
            if body is None:
                continue
            self._try_deflatten_body(body, parent)

    def _try_deflatten_body(self, body: list[Statement], parent: Node):
        """
        Dissolve the state machines of one body as one batch, decided against the tree as it stands
        (`refinery.lib.scripts.ps1.deobfuscation.removal.Ps1RemovalPlan`): the models are read
        once, and only for a body that holds a machine, so a body of many machines costs one build.

        A machine is its dispatcher loop and the store that seeds it, and the two stand or fall
        together: a vetoed half would leave the machine partly dissolved. A machine claiming a
        seed another machine of the batch claims is left for the next run, and so is one whose
        state variable a later statement may read, since the recovered statements store nothing
        into it. A body other than the script's own may run again, and there a statement before the
        seed may read the state the machine left on the run before, unless a constant store into the
        variable comes first.
        """
        cache: Ps1ModelCache | None = None
        plan: Ps1RemovalPlan | None = None
        reads: _StateReads | None = None
        i = 0
        while i < len(body):
            stmt = body[i]
            if not isinstance(stmt, Ps1WhileLoop):
                i += 1
                continue
            match = _match_dispatcher(stmt)
            if match is None:
                i += 1
                continue
            if cache is None:
                cache = model_cache(self, parent)
            name = match.state_var_name
            scope = match.state_var_scope
            init = _find_state_init(body, i, name, scope, cache.call_graph)
            if init is None:
                i += 1
                continue
            init_index, entry_state = init
            if reads is None:
                reads = _StateReads(body, cache.call_graph, reads_in_bodies(tree_root(parent)))
            if reads.after(i, name):
                i += 1
                continue
            first = reads.first(name)
            if (
                not isinstance(parent, Ps1Script)
                and first is not None
                and first < init_index
                and _is_state_assignment(body[first], name, scope) is None
            ):
                i += 1
                continue
            is_exit = _make_exit_check(match.condition, name)
            machine = _extract_state_machine(match, is_exit)
            if machine is None:
                i += 1
                continue
            if entry_state not in machine:
                i += 1
                continue
            if not _exit_condition_is_total(match.condition, name, machine):
                i += 1
                continue
            recovered = _recover_structure(machine, entry_state, is_exit)
            if recovered is None:
                i += 1
                continue
            recovered = [s for s in recovered if _is_state_assignment(s, name, scope) is None]
            # Statements that sit between the state initialization and the dispatcher loop are
            # unrelated and stay where they are; only the loop and its `$state = ...` seed move.
            if plan is None:
                plan = Ps1RemovalPlan(
                    parent,
                    handlers=cache.handlers,
                    under_the_world=False,
                    error_state=cache.error_state,
                )
            plan.propose_together([(body[init_index], None), (stmt, recovered)])
            i += 1
        if plan is not None and plan.commit():
            self.mark_changed()