Module refinery.lib.scripts.vba.deobfuscation.helpers

Shared AST utilities for VBA deobfuscation transforms.

"""
Shared AST utilities for VBA deobfuscation transforms.
"""
from __future__ import annotations

from operator import itemgetter
from typing import Generator

from refinery.lib.scripts import (
    BodyEdit,
    Expression,
    Kind,
    Statement,
    _classify_fields,
    owning_list,
)
from refinery.lib.scripts.vba.deobfuscation.names import CHR_NAMES, CompareMode, Value
from refinery.lib.scripts.vba.model import (
    VbaBinaryExpression,
    VbaBooleanLiteral,
    VbaCallExpression,
    VbaConstDeclaration,
    VbaConstDeclarator,
    VbaEmptyLiteral,
    VbaExpressionStatement,
    VbaFloatLiteral,
    VbaForEachStatement,
    VbaForStatement,
    VbaIdentifier,
    VbaIntegerLiteral,
    VbaLetStatement,
    VbaModule,
    VbaOptionStatement,
    VbaStringLiteral,
    VbaUnaryExpression,
)

LITERAL_TYPES = (VbaStringLiteral, VbaIntegerLiteral, VbaFloatLiteral, VbaBooleanLiteral, VbaEmptyLiteral)


def make_string_literal(value: str) -> VbaStringLiteral:
    escaped = value.replace('"', '""')
    raw = F'"{escaped}"'
    return VbaStringLiteral(value=value, raw=raw)


def is_nan_or_inf(value) -> bool:
    return isinstance(value, float) and (value != value or abs(value) == float('inf'))


def module_compare_mode(module: VbaModule) -> CompareMode:
    """
    Return the module's `Option Compare` mode. VBA `Option Compare` is a module-level directive; the
    default (no directive, or `Binary`) is
    `refinery.lib.scripts.vba.deobfuscation.names.CompareMode.BINARY`, `Text` is case-insensitive,
    and `Database` (Access) uses the database's locale-dependent sort order. The latter cannot be
    reproduced statically, so the folding transforms refuse to fold comparisons under it.
    """
    for stmt in module.body:
        if isinstance(stmt, VbaOptionStatement) and stmt.keyword.lower() == 'compare':
            value = stmt.value.lower()
            if value == 'text':
                return CompareMode.TEXT
            if value == 'database':
                return CompareMode.DATABASE
            return CompareMode.BINARY
    return CompareMode.BINARY


def vba_int_div(a: int | float, b: int | float) -> int:
    """
    VBA integer division (the `\\` operator): both operands are rounded to integers and the
    quotient is truncated toward zero, unlike Python `//` which floors. Raises `ZeroDivisionError`
    when the divisor rounds to zero.
    """
    a, b = round(a), round(b)
    if b == 0:
        raise ZeroDivisionError
    quotient = abs(a) // abs(b)
    return -quotient if (a < 0) != (b < 0) else quotient


def vba_mod(a: int | float, b: int | float) -> int:
    """
    VBA modulo (the `Mod` operator): both operands are rounded to integers and the remainder takes
    the sign of the dividend, unlike Python `%` whose result takes the sign of the divisor. Raises
    `ZeroDivisionError` when the divisor rounds to zero.
    """
    a, b = round(a), round(b)
    if b == 0:
        raise ZeroDivisionError
    remainder = abs(a) % abs(b)
    return -remainder if a < 0 else remainder


def make_integer_literal(value: int) -> VbaIntegerLiteral:
    return VbaIntegerLiteral(value=value, raw=str(value))


def make_float_literal(value: float) -> VbaFloatLiteral:
    return VbaFloatLiteral(value=value, raw=str(value))


def make_numeric_literal(value: int | float) -> VbaIntegerLiteral | VbaFloatLiteral:
    if isinstance(value, float):
        if value == int(value) and abs(value) < 2 ** 53:
            return make_integer_literal(int(value))
        return make_float_literal(value)
    return make_integer_literal(value)


def is_literal(node: Expression) -> bool:
    return isinstance(node, LITERAL_TYPES)


def is_constant_expr(node: Expression) -> bool:
    """
    Returns `True` for expressions that can be safely propagated as constants: literals, Chr/ChrW
    calls with literal integer arguments, and concatenations of such expressions.
    """
    if is_literal(node):
        return True
    if isinstance(node, VbaCallExpression):
        if (
            isinstance(node.callee, VbaIdentifier)
            and node.callee.name.lower() in CHR_NAMES
            and len(node.arguments) == 1
            and node.arguments[0] is not None
            and isinstance(node.arguments[0], VbaIntegerLiteral)
        ):
            return True
        return False
    if isinstance(node, VbaBinaryExpression):
        if node.operator in ('&', '+'):
            return (
                node.left is not None
                and node.right is not None
                and is_constant_expr(node.left)
                and is_constant_expr(node.right)
            )
    if isinstance(node, VbaUnaryExpression):
        if node.operator in ('-', 'Not') and node.operand is not None:
            return is_constant_expr(node.operand)
    return False


def is_identifier_read(node: VbaIdentifier) -> bool:
    """
    Return whether an identifier node is in a read position. Returns `False` for identifiers that
    appear as assignment targets, declaration names, call targets, or loop variables.
    """
    parent = node.parent
    if isinstance(parent, VbaLetStatement) and parent.target is node:
        return False
    if isinstance(parent, (VbaConstDeclaration, VbaConstDeclarator)):
        return False
    if isinstance(parent, VbaCallExpression) and parent.callee is node:
        return False
    if isinstance(parent, VbaExpressionStatement) and parent.expression is node:
        return False
    if (
        isinstance(parent, (VbaForStatement, VbaForEachStatement))
        and parent.variable is node
    ):
        return False
    return True


def literal_value(node: Expression) -> Value:
    if isinstance(node, LITERAL_TYPES):
        return node.value
    return None


def constant_args(arguments: list[Expression | None]) -> list[Value] | None:
    """
    Collect the constant values of a builtin call's arguments, preserving an omitted argument as
    `None`. Returns `None` if any argument is a non-literal expression that cannot be folded.
    """
    values: list[Value] = []
    for arg in arguments:
        if arg is None:
            values.append(None)
        elif is_literal(arg):
            values.append(literal_value(arg))
        else:
            return None
    return values


def string_value(node: Expression | None) -> str | None:
    if isinstance(node, VbaStringLiteral):
        return node.value
    return None


def numeric_value(node: Expression | None) -> int | float | None:
    if isinstance(node, VbaIntegerLiteral):
        return node.value
    if isinstance(node, VbaFloatLiteral):
        return node.value
    return None


def make_chr_call(code_point: int) -> VbaCallExpression:
    return VbaCallExpression(
        callee=VbaIdentifier(name='Chr'),
        arguments=[make_integer_literal(code_point)],
    )


def string_to_expr(value: str) -> Expression:
    """
    Convert a Python string to a VBA AST expression. Printable-only strings become a single string
    literal; strings with non-printable characters become concatenated expressions using Chr calls.
    """
    if not value:
        return make_string_literal('')
    if all(c.isprintable() for c in value):
        return make_string_literal(value)
    parts: list[Expression] = []
    run: list[str] = []
    for c in value:
        if c.isprintable():
            run.append(c)
        else:
            if run:
                parts.append(make_string_literal(''.join(run)))
                run.clear()
            parts.append(make_chr_call(ord(c)))
    if run:
        parts.append(make_string_literal(''.join(run)))
    result = parts[0]
    for part in parts[1:]:
        result = VbaBinaryExpression(left=result, operator='&', right=part)
    return result


def value_to_node(value: Value) -> Expression:
    if value is None:
        return VbaEmptyLiteral()
    if isinstance(value, bool):
        return VbaBooleanLiteral(value=value)
    if isinstance(value, str):
        return string_to_expr(value)
    if isinstance(value, int):
        return make_integer_literal(value)
    return make_numeric_literal(value)


def body_lists(module: VbaModule) -> Generator[list[Statement]]:
    """
    Yield every statement-list body reachable from the module.
    """
    for node in module.walk():
        for field_name, kind in _classify_fields(type(node)):
            if kind != Kind.ChildList:
                continue
            body = getattr(node, field_name)
            if body and isinstance(body[0], Statement):
                yield body


def apply_removals(removals: list[tuple[int, list[Statement]]]) -> bool:
    """
    Delete the statements at the given (body, index) positions through `BodyEdit`, so every removal
    advances the mutation counter. Indices are resolved in reverse order — a deletion must not
    shift a later one — and each statement is then spliced by identity. Returns whether any
    removals occurred.
    """
    if not removals:
        return False
    removals.sort(key=itemgetter(0), reverse=True)
    changed = False
    for pos, body in removals:
        node = body[pos]
        holder = owning_list(node)
        if holder is None:
            continue
        parent, attr = holder
        edit = BodyEdit(parent, attr)
        edit.splice(node, [])
        if edit.apply():
            changed = True
    return changed

Functions

def make_string_literal(value)

Expand source code Browse git

def make_string_literal(value: str) -> VbaStringLiteral:
    escaped = value.replace('"', '""')
    raw = F'"{escaped}"'
    return VbaStringLiteral(value=value, raw=raw)

def is_nan_or_inf(value)

Expand source code Browse git

def is_nan_or_inf(value) -> bool:
    return isinstance(value, float) and (value != value or abs(value) == float('inf'))

def module_compare_mode(module)

Return the module's Option Compare mode. VBA Option Compare is a module-level directive; the default (no directive, or Binary) is CompareMode.BINARY, Text is case-insensitive, and Database (Access) uses the database's locale-dependent sort order. The latter cannot be reproduced statically, so the folding transforms refuse to fold comparisons under it.

Expand source code Browse git

def module_compare_mode(module: VbaModule) -> CompareMode:
    """
    Return the module's `Option Compare` mode. VBA `Option Compare` is a module-level directive; the
    default (no directive, or `Binary`) is
    `refinery.lib.scripts.vba.deobfuscation.names.CompareMode.BINARY`, `Text` is case-insensitive,
    and `Database` (Access) uses the database's locale-dependent sort order. The latter cannot be
    reproduced statically, so the folding transforms refuse to fold comparisons under it.
    """
    for stmt in module.body:
        if isinstance(stmt, VbaOptionStatement) and stmt.keyword.lower() == 'compare':
            value = stmt.value.lower()
            if value == 'text':
                return CompareMode.TEXT
            if value == 'database':
                return CompareMode.DATABASE
            return CompareMode.BINARY
    return CompareMode.BINARY

def vba_int_div(a, b)

VBA integer division (the \ operator): both operands are rounded to integers and the quotient is truncated toward zero, unlike Python // which floors. Raises ZeroDivisionError when the divisor rounds to zero.

Expand source code Browse git

def vba_int_div(a: int | float, b: int | float) -> int:
    """
    VBA integer division (the `\\` operator): both operands are rounded to integers and the
    quotient is truncated toward zero, unlike Python `//` which floors. Raises `ZeroDivisionError`
    when the divisor rounds to zero.
    """
    a, b = round(a), round(b)
    if b == 0:
        raise ZeroDivisionError
    quotient = abs(a) // abs(b)
    return -quotient if (a < 0) != (b < 0) else quotient

def vba_mod(a, b)

VBA modulo (the Mod operator): both operands are rounded to integers and the remainder takes the sign of the dividend, unlike Python % whose result takes the sign of the divisor. Raises ZeroDivisionError when the divisor rounds to zero.

Expand source code Browse git

def vba_mod(a: int | float, b: int | float) -> int:
    """
    VBA modulo (the `Mod` operator): both operands are rounded to integers and the remainder takes
    the sign of the dividend, unlike Python `%` whose result takes the sign of the divisor. Raises
    `ZeroDivisionError` when the divisor rounds to zero.
    """
    a, b = round(a), round(b)
    if b == 0:
        raise ZeroDivisionError
    remainder = abs(a) % abs(b)
    return -remainder if a < 0 else remainder

def make_integer_literal(value)

Expand source code Browse git

def make_integer_literal(value: int) -> VbaIntegerLiteral:
    return VbaIntegerLiteral(value=value, raw=str(value))

def make_float_literal(value)

Expand source code Browse git

def make_float_literal(value: float) -> VbaFloatLiteral:
    return VbaFloatLiteral(value=value, raw=str(value))

def make_numeric_literal(value)

Expand source code Browse git

def make_numeric_literal(value: int | float) -> VbaIntegerLiteral | VbaFloatLiteral:
    if isinstance(value, float):
        if value == int(value) and abs(value) < 2 ** 53:
            return make_integer_literal(int(value))
        return make_float_literal(value)
    return make_integer_literal(value)

def is_literal(node)

Expand source code Browse git

def is_literal(node: Expression) -> bool:
    return isinstance(node, LITERAL_TYPES)

def is_constant_expr(node)

Returns True for expressions that can be safely propagated as constants: literals, Chr/ChrW calls with literal integer arguments, and concatenations of such expressions.

Expand source code Browse git

def is_constant_expr(node: Expression) -> bool:
    """
    Returns `True` for expressions that can be safely propagated as constants: literals, Chr/ChrW
    calls with literal integer arguments, and concatenations of such expressions.
    """
    if is_literal(node):
        return True
    if isinstance(node, VbaCallExpression):
        if (
            isinstance(node.callee, VbaIdentifier)
            and node.callee.name.lower() in CHR_NAMES
            and len(node.arguments) == 1
            and node.arguments[0] is not None
            and isinstance(node.arguments[0], VbaIntegerLiteral)
        ):
            return True
        return False
    if isinstance(node, VbaBinaryExpression):
        if node.operator in ('&', '+'):
            return (
                node.left is not None
                and node.right is not None
                and is_constant_expr(node.left)
                and is_constant_expr(node.right)
            )
    if isinstance(node, VbaUnaryExpression):
        if node.operator in ('-', 'Not') and node.operand is not None:
            return is_constant_expr(node.operand)
    return False

def is_identifier_read(node)

Return whether an identifier node is in a read position. Returns False for identifiers that appear as assignment targets, declaration names, call targets, or loop variables.

Expand source code Browse git

def is_identifier_read(node: VbaIdentifier) -> bool:
    """
    Return whether an identifier node is in a read position. Returns `False` for identifiers that
    appear as assignment targets, declaration names, call targets, or loop variables.
    """
    parent = node.parent
    if isinstance(parent, VbaLetStatement) and parent.target is node:
        return False
    if isinstance(parent, (VbaConstDeclaration, VbaConstDeclarator)):
        return False
    if isinstance(parent, VbaCallExpression) and parent.callee is node:
        return False
    if isinstance(parent, VbaExpressionStatement) and parent.expression is node:
        return False
    if (
        isinstance(parent, (VbaForStatement, VbaForEachStatement))
        and parent.variable is node
    ):
        return False
    return True

def literal_value(node)

Expand source code Browse git

def literal_value(node: Expression) -> Value:
    if isinstance(node, LITERAL_TYPES):
        return node.value
    return None

def constant_args(arguments)

Collect the constant values of a builtin call's arguments, preserving an omitted argument as None. Returns None if any argument is a non-literal expression that cannot be folded.

Expand source code Browse git

def constant_args(arguments: list[Expression | None]) -> list[Value] | None:
    """
    Collect the constant values of a builtin call's arguments, preserving an omitted argument as
    `None`. Returns `None` if any argument is a non-literal expression that cannot be folded.
    """
    values: list[Value] = []
    for arg in arguments:
        if arg is None:
            values.append(None)
        elif is_literal(arg):
            values.append(literal_value(arg))
        else:
            return None
    return values

def string_value(node)

Expand source code Browse git

def string_value(node: Expression | None) -> str | None:
    if isinstance(node, VbaStringLiteral):
        return node.value
    return None

def numeric_value(node)

Expand source code Browse git

def numeric_value(node: Expression | None) -> int | float | None:
    if isinstance(node, VbaIntegerLiteral):
        return node.value
    if isinstance(node, VbaFloatLiteral):
        return node.value
    return None

def make_chr_call(code_point)

Expand source code Browse git

def make_chr_call(code_point: int) -> VbaCallExpression:
    return VbaCallExpression(
        callee=VbaIdentifier(name='Chr'),
        arguments=[make_integer_literal(code_point)],
    )

def string_to_expr(value)

Convert a Python string to a VBA AST expression. Printable-only strings become a single string literal; strings with non-printable characters become concatenated expressions using Chr calls.

Expand source code Browse git

def string_to_expr(value: str) -> Expression:
    """
    Convert a Python string to a VBA AST expression. Printable-only strings become a single string
    literal; strings with non-printable characters become concatenated expressions using Chr calls.
    """
    if not value:
        return make_string_literal('')
    if all(c.isprintable() for c in value):
        return make_string_literal(value)
    parts: list[Expression] = []
    run: list[str] = []
    for c in value:
        if c.isprintable():
            run.append(c)
        else:
            if run:
                parts.append(make_string_literal(''.join(run)))
                run.clear()
            parts.append(make_chr_call(ord(c)))
    if run:
        parts.append(make_string_literal(''.join(run)))
    result = parts[0]
    for part in parts[1:]:
        result = VbaBinaryExpression(left=result, operator='&', right=part)
    return result

def value_to_node(value)

Expand source code Browse git

def value_to_node(value: Value) -> Expression:
    if value is None:
        return VbaEmptyLiteral()
    if isinstance(value, bool):
        return VbaBooleanLiteral(value=value)
    if isinstance(value, str):
        return string_to_expr(value)
    if isinstance(value, int):
        return make_integer_literal(value)
    return make_numeric_literal(value)

def body_lists(module)

Yield every statement-list body reachable from the module.

Expand source code Browse git

def body_lists(module: VbaModule) -> Generator[list[Statement]]:
    """
    Yield every statement-list body reachable from the module.
    """
    for node in module.walk():
        for field_name, kind in _classify_fields(type(node)):
            if kind != Kind.ChildList:
                continue
            body = getattr(node, field_name)
            if body and isinstance(body[0], Statement):
                yield body

def apply_removals(removals)

Delete the statements at the given (body, index) positions through BodyEdit, so every removal advances the mutation counter. Indices are resolved in reverse order — a deletion must not shift a later one — and each statement is then spliced by identity. Returns whether any removals occurred.

Expand source code Browse git

def apply_removals(removals: list[tuple[int, list[Statement]]]) -> bool:
    """
    Delete the statements at the given (body, index) positions through `BodyEdit`, so every removal
    advances the mutation counter. Indices are resolved in reverse order — a deletion must not
    shift a later one — and each statement is then spliced by identity. Returns whether any
    removals occurred.
    """
    if not removals:
        return False
    removals.sort(key=itemgetter(0), reverse=True)
    changed = False
    for pos, body in removals:
        node = body[pos]
        holder = owning_list(node)
        if holder is None:
            continue
        parent, attr = holder
        edit = BodyEdit(parent, attr)
        edit.splice(node, [])
        if edit.apply():
            changed = True
    return changed