From 603ac5761bb850238a49c85a129e72ed82fae823 Mon Sep 17 00:00:00 2001
From: Takahiro Yoshimura <alterakey@protonmail.com>
Date: Tue, 17 Sep 2024 23:40:15 +0900
Subject: [PATCH] Experimentally adding command to merge slices

---
 trueseeing/app/cmd/android/asm.py  | 33 ++++++++++++++++++++++++++++++
 trueseeing/core/android/asm.py     | 21 +++++++++++++++++++
 trueseeing/core/android/context.py |  4 ++++
 3 files changed, 58 insertions(+)

diff --git a/trueseeing/app/cmd/android/asm.py b/trueseeing/app/cmd/android/asm.py
index 000df67..09bbae9 100644
--- a/trueseeing/app/cmd/android/asm.py
+++ b/trueseeing/app/cmd/android/asm.py
@@ -29,6 +29,8 @@ def get_commands(self) -> CommandMap:
       'cd!':dict(e=self._disassemble, t={'apk'}),
       'cds':dict(e=self._disassemble_nodex, t={'apk'}),
       'cds!':dict(e=self._disassemble_nodex, t={'apk'}),
+      'cm':dict(e=self._merge, n='cm[!]', d='merge slices', t={'xapk'}),
+      'cm!':dict(e=self._merge, t={'xapk'}),
       'co':dict(e=self._export_context, n='co[!] /path [pat]', d='export codebase', t={'apk'}),
       'co!':dict(e=self._export_context, t={'apk'}),
     }
@@ -240,3 +242,34 @@ def _deduce_archive_format(self, path: str) -> ArchiveFormat:
       return 'tar:gz'
     else:
       return None
+
+  async def _merge(self, args: deque[str]) -> None:
+    target = self._helper.require_target('need target')
+
+    cmd = args.popleft()
+
+    import os
+    import time
+    from tempfile import TemporaryDirectory
+    from trueseeing.core.android.asm import APKAssembler
+    from trueseeing.core.android.tools import move_apk
+
+    apk = target.replace('.xapk', '.apk')
+    origapk = apk.replace('.apk', '.apk.orig')
+
+    if os.path.exists(origapk) and not cmd.endswith('!'):
+      ui.fatal('backup file exists; force (!) to overwrite')
+
+    ui.info('merging slices {target} -> {apk}'.format(target=target, apk=apk))
+
+    at = time.time()
+
+    with TemporaryDirectory() as td:
+      outapk, outsig = await APKAssembler.merge_slices(target, td)
+
+      if os.path.exists(apk):
+        move_apk(apk, origapk)
+
+      move_apk(outapk, apk)
+
+    ui.success('done ({t:.02f} sec.)'.format(t=(time.time() - at)))
diff --git a/trueseeing/core/android/asm.py b/trueseeing/core/android/asm.py
index 34a9961..cd35ed0 100644
--- a/trueseeing/core/android/asm.py
+++ b/trueseeing/core/android/asm.py
@@ -142,6 +142,27 @@ async def assemble_from_path(cls, wd: str, path: str) -> Tuple[str, str]:
 
     return os.path.join(wd, 'output.apk'), os.path.join(wd, 'output.apk.idsig')
 
+  @classmethod
+  async def merge_slices(cls, xapk: str, wd: str) -> Tuple[str, str]:
+    import os
+    from trueseeing.core.tools import invoke_streaming
+    from trueseeing.core.android.tools import toolchains
+
+    pub.sendMessage('progress.core.asm.asm.begin')
+
+    with toolchains() as tc:
+      async for l in invoke_streaming(
+        'java -jar {apkeditor} m -i {xapk} -o {wd}/output.apk'.format(
+          wd=wd, xapk=shlex.quote(xapk),
+          apkeditor=tc['apkeditor'],
+        ), redir_stderr=True
+      ):
+        pub.sendMessage('progress.core.asm.asm.update')
+
+    pub.sendMessage('progress.core.asm.asm.done')
+
+    return os.path.join(wd, 'output.apk'), os.path.join(wd, 'output.apk.idsig')
+
 class SigningKey:
   _path: str
 
diff --git a/trueseeing/core/android/context.py b/trueseeing/core/android/context.py
index 5f30765..37b7d5a 100644
--- a/trueseeing/core/android/context.py
+++ b/trueseeing/core/android/context.py
@@ -379,6 +379,7 @@ def is_qualname_excluded(self, qualname: Optional[str]) -> bool:
 
 class XAPKContext(APKContext):
   _disasm: Optional[APKDisassembler] = None
+  _type: ClassVar[Set[ContextType]] = {'xapk'}
 
   async def _get_disassembler(self) -> APKDisassembler:
     assert self._disasm
@@ -404,6 +405,9 @@ async def _analyze(self, level: int) -> None:
       finally:
         self._disasm = None
 
+  def _get_type(self) -> Set[ContextType]:
+    return super()._type | self._type
+
   async def _get_info(self, extended: bool) -> AsyncIterator[ContextInfo]:
     async for m in super()._get_info(extended):
       yield m