split-dex 0.1.0__py3-none-any.whl

split_dex/__init__.py

@@ -0,0 +1,5 @@
+"""split-dex: split the classes of an Android APK into grouped DEX files."""
+
+
+class Error(Exception):
+    """Base error for split-dex."""

split_dex/class_sets.py

@@ -0,0 +1,135 @@
+"""Sort the classes of an APK into named sets, by patterns on their names.
+
+Each set is grouped by its own strategy, and two sets never share a DEX file.
+
+Patterns match binary class names, such as `com.myapp.Main` or `com.myapp.Outer$Inner`. In a
+pattern, `*` matches any characters except a dot, and `**` matches any characters, dots
+included. So `com.myapp.**` matches every class under `com.myapp`, `com.myapp.*` matches the
+classes directly in `com.myapp`, and `**` matches every class.
+"""
+
+import collections.abc
+import dataclasses
+import re
+
+import split_dex
+from split_dex import strategy
+
+# How many unmatched classes an error message names.
+_EXAMPLES = 3
+
+
+class AssignmentError(split_dex.Error):
+    """Raised when the classes cannot be sorted into the sets."""
+
+
+@dataclasses.dataclass(frozen=True)
+class ClassSet:
+    """A named set of classes, and the strategy that groups them.
+
+    Attributes:
+        name: The set's name, written to sets.csv next to each of its DEX files.
+        strategy: How the classes of the set are grouped.
+    """
+
+    name: str
+    strategy: strategy.Strategy
+
+
+@dataclasses.dataclass(frozen=True)
+class Assignment:
+    """Sends the classes whose names match `pattern` to the set named `set_name`.
+
+    Attributes:
+        pattern: A pattern on binary class names, such as `com.myapp.**`.
+        set_name: The name of a set.
+    """
+
+    pattern: str
+    set_name: str
+
+
+def binary_name(descriptor: str) -> str:
+    """The binary name of a class: `Lcom/a/B$C;` gives `com.a.B$C`.
+
+    Args:
+        descriptor: A class descriptor.
+
+    Returns:
+        The binary name.
+    """
+    return descriptor[1:-1].replace("/", ".")
+
+
+def assign(
+    descriptors: collections.abc.Sequence[str],
+    sets: collections.abc.Sequence[ClassSet],
+    assignments: collections.abc.Sequence[Assignment],
+) -> dict[str, list[int]]:
+    """Sort classes into sets.
+
+    Each class goes to the set of the first assignment whose pattern matches its name. With no
+    assignments and a single set, every class goes to that set.
+
+    Args:
+        descriptors: The descriptor of every class, by class id.
+        sets: The sets.
+        assignments: The assignments, in the order they are tried.
+
+    Returns:
+        The ids of the classes of each set, in class id order, by set name. A set that no class
+        matches has an empty list.
+
+    Raises:
+        AssignmentError: If two sets have the same name, an assignment names an unknown set,
+            there are several sets but no assignments, or a class matches no pattern.
+    """
+    members: dict[str, list[int]] = {}
+    for class_set in sets:
+        if class_set.name in members:
+            raise AssignmentError(f"two sets are named {class_set.name!r}")
+        members[class_set.name] = []
+    if len(assignments) == 0:
+        if len(sets) != 1:
+            raise AssignmentError("with several sets, assignments must say which classes go where")
+        members[sets[0].name] = list(range(len(descriptors)))
+        return members
+    rules = [(_compile(rule.pattern), _members_of(members, rule.set_name)) for rule in assignments]
+    unmatched: list[str] = []
+    for class_id, descriptor in enumerate(descriptors):
+        name = binary_name(descriptor)
+        target = next((ids for pattern, ids in rules if pattern.fullmatch(name) is not None), None)
+        if target is None:
+            unmatched.append(name)
+        else:
+            target.append(class_id)
+    if len(unmatched) > 0:
+        raise AssignmentError(
+            f"{len(unmatched)} classes match no pattern, such as "
+            f"{', '.join(unmatched[:_EXAMPLES])}; add a last pattern such as '**' to catch them"
+        )
+    return members
+
+
+def _members_of(members: dict[str, list[int]], set_name: str) -> list[int]:
+    ids = members.get(set_name)
+    if ids is None:
+        raise AssignmentError(f"no set is named {set_name!r}")
+    return ids
+
+
+def _compile(pattern: str) -> re.Pattern[str]:
+    """The regular expression of a pattern: `**` is any text, `*` any text without a dot."""
+    parts: list[str] = []
+    index = 0
+    while index < len(pattern):
+        if pattern.startswith("**", index):
+            parts.append(".*")
+            index += 2
+        elif pattern[index] == "*":
+            parts.append("[^.]*")
+            index += 1
+        else:
+            parts.append(re.escape(pattern[index]))
+            index += 1
+    return re.compile("".join(parts))

split_dex/cli.py

@@ -0,0 +1,227 @@
+"""The split-dex command line."""
+
+import argparse
+import collections.abc
+import importlib.metadata
+import pathlib
+import re
+import sys
+from typing import Any
+
+import split_dex
+from split_dex import class_sets
+from split_dex import registry
+from split_dex import splitter
+from split_dex import strategy
+
+_EPILOG = """\
+sets:
+  Each class goes to the set of the first --assign pattern that matches its
+  name, and each set is split by its own strategy. Two sets never share a DEX
+  file. Without --set, every class is in one set, "all", split by
+  package-affinity. Patterns match class names such as com.myapp.Main: "*"
+  matches any characters except a dot, "**" any characters. "@FILE" reads
+  patterns from a file, one per line.
+
+strategies:
+  STRATEGY is a built-in or plugin name (see --list-strategies), or your own
+  strategy class as MODULE:CLASS or FILE.py:CLASS. Its settings are the
+  keyword parameters of the class's constructor.
+
+examples:
+  split-dex --apk app.apk --out out/
+  split-dex --apk app.apk --out out/ --set all=package-affinity,cap=500
+  split-dex --apk app.apk --out out/ \\
+      --set app=package-affinity,cap=500 --set libraries=single \\
+      --assign 'com.myapp.**=app' --assign '**=libraries'
+  split-dex --apk app.apk --out out/ --set all=./my_strategy.py:MyStrategy,k=8
+
+output in DIR:
+  <k>.dex     the grouped DEX files: 0.dex, 1.dex, ..., set by set
+  index.csv   class,file: every class and the DEX file it is in
+  sets.csv    file,set: the set of each DEX file
+"""
+
+_BYTES_PER_MB = 1048576
+
+# A set name is written to sets.csv, so it cannot hold a comma or a line break.
+_SET_NAME = re.compile(r"[A-Za-z0-9_.-]+")
+
+
+class _UsageError(split_dex.Error):
+    """Raised when an option's value is malformed."""
+
+
+def main(argv: collections.abc.Sequence[str] | None = None) -> int:
+    """Run split-dex with command-line arguments.
+
+    Prints a summary of the split on stdout, and a warning on stderr for each block, such as a
+    package, that was too big for one DEX file.
+
+    Args:
+        argv: The arguments, without the program name. None reads them from `sys.argv`.
+
+    Returns:
+        The exit status: 0 when the split succeeds, 1 when it fails. Bad arguments exit with
+        status 2 instead, through argparse's SystemExit.
+    """
+    parser = _parser()
+    arguments = parser.parse_args(argv)
+    try:
+        options = _options(arguments)
+    except (_UsageError, strategy.StrategyError) as error:
+        parser.error(str(error))
+    try:
+        result = splitter.split(options)
+    except split_dex.Error as error:
+        print(f"split-dex: {error}", file=sys.stderr)
+        return 1
+    _print_summary(result)
+    return 0
+
+
+class _ListStrategies(argparse.Action):
+    """Prints the built-in and plugin strategies and exits, like --help."""
+
+    def __call__(
+        self,
+        parser: argparse.ArgumentParser,
+        namespace: argparse.Namespace,
+        values: Any,
+        option_string: str | None = None,
+    ) -> None:
+        try:
+            lines = registry.describe()
+        except strategy.StrategyError as error:
+            parser.exit(1, f"split-dex: {error}\n")
+        print("\n".join(lines))
+        parser.exit()
+
+
+def _parser() -> argparse.ArgumentParser:
+    version = importlib.metadata.version("split-dex")
+    parser = argparse.ArgumentParser(
+        prog="split-dex",
+        description=f"split-dex {version}: split an APK's classes into grouped DEX files.",
+        epilog=_EPILOG,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    required = parser.add_argument_group("required arguments")
+    required.add_argument(
+        "--apk",
+        required=True,
+        type=_non_empty,
+        metavar="FILE",
+        help="input APK; all its classes*.dex entries are read",
+    )
+    required.add_argument(
+        "--out", required=True, type=_non_empty, metavar="DIR", help="output directory"
+    )
+    parser.add_argument(
+        "--set",
+        action="append",
+        default=[],
+        dest="sets",
+        metavar="NAME=STRATEGY[,KEY=VALUE...]",
+        help="a set of classes and the strategy that splits it, with the strategy's settings; "
+        "repeat for more sets (DEX files are written in this order)",
+    )
+    parser.add_argument(
+        "--assign",
+        action="append",
+        default=[],
+        dest="assignments",
+        metavar="PATTERN[,PATTERN...]=NAME",
+        help="send the classes whose names match a pattern to the set NAME; repeat, the first "
+        "match wins",
+    )
+    parser.add_argument(
+        "--list-strategies",
+        action=_ListStrategies,
+        nargs=0,
+        help="show the strategies and their settings, and exit",
+    )
+    return parser
+
+
+def _non_empty(value: str) -> str:
+    if value == "":
+        raise argparse.ArgumentTypeError("must not be empty")
+    return value
+
+
+def _options(arguments: argparse.Namespace) -> splitter.Options:
+    sets = [_parse_set(text) for text in arguments.sets]
+    if len(sets) == 0:
+        sets = list(splitter.default_sets())
+    names = [class_set.name for class_set in sets]
+    if len(set(names)) < len(names):
+        raise _UsageError("two --set options have the same name")
+    assignments = [rule for text in arguments.assignments for rule in _parse_assignment(text)]
+    unknown = sorted({rule.set_name for rule in assignments} - set(names))
+    if len(unknown) > 0:
+        raise _UsageError(f"--assign names a set that no --set declares: {', '.join(unknown)}")
+    if len(sets) > 1 and len(assignments) == 0:
+        raise _UsageError("with several --set options, --assign must say which classes go where")
+    return splitter.Options(
+        apk_path=pathlib.Path(arguments.apk),
+        out_dir=pathlib.Path(arguments.out),
+        sets=sets,
+        assignments=assignments,
+    )
+
+
+def _parse_set(text: str) -> class_sets.ClassSet:
+    """A set from `NAME=STRATEGY[,KEY=VALUE...]`."""
+    name, separator, spec = text.partition("=")
+    if separator == "" or spec == "":
+        raise _UsageError(f"--set {text}: expected NAME=STRATEGY[,KEY=VALUE...]")
+    if _SET_NAME.fullmatch(name) is None:
+        raise _UsageError(f"--set {text}: a set name uses only letters, digits, '_', '.' and '-'")
+    strategy_name, *setting_texts = spec.split(",")
+    settings: dict[str, str] = {}
+    for setting in setting_texts:
+        key, separator, value = setting.partition("=")
+        if separator == "" or key == "":
+            raise _UsageError(f"--set {text}: expected KEY=VALUE, not {setting!r}")
+        settings[key] = value
+    return class_sets.ClassSet(name, registry.create(strategy_name, settings))
+
+
+def _parse_assignment(text: str) -> list[class_sets.Assignment]:
+    """Assignments from `PATTERN[,PATTERN...]=NAME`, where `@FILE` reads patterns from a file."""
+    patterns_text, separator, set_name = text.rpartition("=")
+    if separator == "" or patterns_text == "" or set_name == "":
+        raise _UsageError(f"--assign {text}: expected PATTERN[,PATTERN...]=NAME")
+    patterns: list[str] = []
+    for pattern in patterns_text.split(","):
+        if pattern.startswith("@"):
+            patterns.extend(_read_patterns(pathlib.Path(pattern[1:])))
+        else:
+            patterns.append(pattern)
+    return [class_sets.Assignment(pattern, set_name) for pattern in patterns]
+
+
+def _read_patterns(path: pathlib.Path) -> list[str]:
+    """The patterns of a file: one per line. Blank lines and lines starting with '#' are skipped."""
+    try:
+        lines = path.read_bytes().decode("utf-8").split("\n")
+    except (OSError, UnicodeDecodeError) as error:
+        raise _UsageError(f"cannot read patterns from {path}: {error}") from error
+    patterns = [line.strip() for line in lines]
+    return [pattern for pattern in patterns if pattern != "" and pattern.startswith("#") is False]
+
+
+def _print_summary(result: splitter.Result) -> None:
+    print(f"total_classes={result.total_classes} duplicate_classes={result.duplicate_classes}")
+    print(
+        f"total_dex={result.dex_files} overflow_splits={result.overflow_splits} "
+        f"total_mb={result.total_bytes / _BYTES_PER_MB:.1f}"
+    )
+    for set_result in result.sets:
+        print(f"set={set_result.name} classes={set_result.classes} dex={set_result.dex_files}")
+    for block in result.split_blocks:
+        print(
+            f"split-dex: warning: {block} is too big for one DEX file and was split across files",
+            file=sys.stderr,
+        )

split_dex/partition.py

@@ -0,0 +1,229 @@
+"""Group classes by how they are linked, without splitting a Java package.
+
+A class usually loads the classes it is linked to, so linked classes should land in the same DEX
+file. A package is never split across two groups: when each DEX file is served by its own class
+loader, two classes of one package in different loaders cannot use each other's package-private
+members, and ART throws IllegalAccessError. Keeping each package whole keeps that access inside
+one class loader.
+"""
+
+import collections.abc
+import dataclasses
+import heapq
+
+
+class AffinityGraph:
+    """An undirected, weighted graph of the links between classes.
+
+    Each class is a dense id in `range(size)`. The weight of the edge between two classes counts
+    the links between them. The graph is symmetric: an edge is stored on both of its ends.
+    """
+
+    def __init__(self, size: int) -> None:
+        """Create a graph of `size` classes with no edges.
+
+        Args:
+            size: The number of classes.
+        """
+        self._neighbours: list[dict[int, int]] = [{} for _ in range(size)]
+
+    @property
+    def size(self) -> int:
+        """The number of classes in the graph."""
+        return len(self._neighbours)
+
+    def neighbours(self, class_id: int) -> collections.abc.Mapping[int, int]:
+        """The classes linked to a class.
+
+        Args:
+            class_id: The class.
+
+        Returns:
+            A mapping from each linked class to the weight of its edge.
+        """
+        return self._neighbours[class_id]
+
+    def add_edge(self, first: int, second: int, weight: int = 1) -> None:
+        """Add `weight` to the edge between two classes, on both of its ends.
+
+        An edge from a class to itself is ignored.
+
+        Args:
+            first: One end of the edge.
+            second: The other end of the edge.
+            weight: How much to add to the weight of the edge.
+        """
+        if first == second:
+            return
+        self._neighbours[first][second] = self._neighbours[first].get(second, 0) + weight
+        self._neighbours[second][first] = self._neighbours[second].get(first, 0) + weight
+
+
+def package_of(descriptor: str) -> str:
+    """The package of a class descriptor: the descriptor up to its last '/'.
+
+    For example `Lcom/a/B;` gives `Lcom/a`. A class in the unnamed package, such as `LFoo;`,
+    gives `L`.
+
+    Args:
+        descriptor: A class descriptor.
+
+    Returns:
+        The package, in descriptor form.
+    """
+    slash = descriptor.rfind("/")
+    if slash == -1:
+        return "L"
+    return descriptor[:slash]
+
+
+def partition_by_package_affinity(
+    members: collections.abc.Sequence[int],
+    graph: AffinityGraph,
+    descriptors: collections.abc.Sequence[str],
+    cap: int,
+) -> list[list[int]]:
+    """Partition classes into groups of about `cap` classes, keeping each package whole.
+
+    The steps:
+
+    1. Group the members by package. A package is the unit that moves; it is never split.
+    2. Add up the edges between classes into edges between packages.
+    3. Put the packages in order so that linked packages are next to each other. The order
+       starts from the package linked to the most other packages. Each next package is the one
+       with the heaviest edge to the packages already placed.
+    4. Fill groups of about `cap` classes in that order, each package whole. A package of `cap`
+       classes or more ends the group it is in, so a package bigger than `cap` gets a group of
+       its own.
+
+    Packages with no links come last and are packed together, so no group is left with a single
+    stray class. The result is deterministic: ties are broken by package ids, which follow the
+    order in which `members` first shows each package, never by hash order.
+
+    Args:
+        members: The ids of the classes to partition, a subset of the classes of `graph`.
+        graph: The links between all classes. Only edges between two members are used.
+        descriptors: The descriptor of every class, indexed by class id.
+        cap: The target number of classes per group. A value below 1 counts as 1.
+
+    Returns:
+        The groups, each a list of class ids. No group is empty.
+    """
+    packages = _group_by_package(members, descriptors)
+    edges = _package_edges(members, graph, packages)
+    order = _order_packages(edges)
+    return fill_groups((packages.classes[package] for package in order), cap)
+
+
+def fill_groups(
+    units: collections.abc.Iterable[collections.abc.Sequence[int]], cap: int
+) -> list[list[int]]:
+    """Fill groups of about `cap` classes with whole units, such as packages, in order.
+
+    A unit that would push the current group over `cap` starts a new group, and a group that
+    reaches `cap` is closed. So a unit of `cap` classes or more gets a group of its own.
+
+    Args:
+        units: The units, each a sequence of class ids, in the order they are filled in.
+        cap: The target number of classes per group. A value below 1 counts as 1.
+
+    Returns:
+        The groups, each a list of class ids. No group is empty.
+    """
+    cap = max(cap, 1)
+    groups: list[list[int]] = []
+    current: list[int] = []
+    for unit in units:
+        if len(current) > 0 and len(current) + len(unit) > cap:
+            groups.append(current)
+            current = []
+        current.extend(unit)
+        if len(current) >= cap:
+            groups.append(current)
+            current = []
+    if len(current) > 0:
+        groups.append(current)
+    return groups
+
+
+@dataclasses.dataclass(frozen=True)
+class _Packages:
+    """The packages of a set of classes. Package ids count from 0, in order of appearance.
+
+    Attributes:
+        classes: The classes of each package, indexed by package id.
+        package_of_class: The package id of each class.
+    """
+
+    classes: list[list[int]]
+    package_of_class: dict[int, int]
+
+
+def _group_by_package(
+    members: collections.abc.Sequence[int], descriptors: collections.abc.Sequence[str]
+) -> _Packages:
+    package_ids: dict[str, int] = {}
+    classes: list[list[int]] = []
+    package_of_class: dict[int, int] = {}
+    for class_id in members:
+        package_id = package_ids.setdefault(package_of(descriptors[class_id]), len(classes))
+        if package_id == len(classes):
+            classes.append([])
+        classes[package_id].append(class_id)
+        package_of_class[class_id] = package_id
+    return _Packages(classes, package_of_class)
+
+
+def _package_edges(
+    members: collections.abc.Sequence[int], graph: AffinityGraph, packages: _Packages
+) -> list[dict[int, int]]:
+    """Add up the edges between classes of two different packages into package edges."""
+    edges: list[dict[int, int]] = [{} for _ in packages.classes]
+    for class_id in members:
+        package = packages.package_of_class[class_id]
+        for neighbour, weight in graph.neighbours(class_id).items():
+            neighbour_package = packages.package_of_class.get(neighbour)
+            if neighbour_package is None or neighbour_package == package:
+                continue
+            edges[package][neighbour_package] = edges[package].get(neighbour_package, 0) + weight
+    return edges
+
+
+def _order_packages(
+    edges: collections.abc.Sequence[collections.abc.Mapping[int, int]],
+) -> list[int]:
+    """Order the packages so that linked packages are next to each other.
+
+    Walks start from the packages with the most links (ties: the lower id first). Each walk is
+    best-first, like Prim's algorithm: it takes the package with the heaviest edge to the
+    packages already placed.
+    """
+    seeds = sorted(range(len(edges)), key=lambda package: (-len(edges[package]), package))
+    placed = [False] * len(edges)
+    order: list[int] = []
+    for seed in seeds:
+        if placed[seed] is False:
+            _walk_from(seed, edges, placed, order)
+    return order
+
+
+def _walk_from(
+    seed: int,
+    edges: collections.abc.Sequence[collections.abc.Mapping[int, int]],
+    placed: list[bool],
+    order: list[int],
+) -> None:
+    """Place every package reachable from `seed`, heaviest edge first, and append it to `order`."""
+    # A max-heap of (edge weight, package id), kept as a min-heap of the negated values: the
+    # heaviest edge comes out first, and of equal weights the higher package id.
+    heap = [(0, -seed)]
+    while len(heap) > 0:
+        _, negated_package = heapq.heappop(heap)
+        package = -negated_package
+        if placed[package] is True:
+            continue
+        placed[package] = True
+        order.append(package)
+        for neighbour, weight in edges[package].items():
+            if placed[neighbour] is False:
+                heapq.heappush(heap, (-weight, -neighbour))