import unittest from pathlib import Path from lihont import media_patch ROOT = Path(__file__).resolve().parents[1] SMALI24 = ROOT / "work" / "tabs" / "smali24" / "X" REAL_SOURCE = SMALI24 / "0qvH.2.smali" REAL_CONTEXT = SMALI24 / "0qwE.2.smali" AWEME = media_patch.AWEME SOURCE = """\ .class public final LX/0qvH; .super Ljava/lang/Object; .method public final LIZ(%(aweme)sZ)V .registers 7 return-void .end method .method public final LIZIZ(%(aweme)s)V .registers 11 .prologue .line 16777216 iget-object v1, p0, LX/0qvH;->LJ:Ljava/lang/String; return-void .end method """ % {"aweme": AWEME} CONTEXT = """\ .class public final LX/0qwE; .super LX/0qzh; .method public final LJFF()Z .registers 2 iget-boolean v0, p0, LX/0qwE;->LJIJI:Z return v0 .end method .method public final LJIIIZ()Z .registers 2 .prologue .line 0 iget-boolean v0, p0, LX/0qwE;->LJIJ:Z .line 1 return v0 .end method .method public final LJIIJ()Z .registers 2 iget-boolean v0, p0, LX/0qwE;->LJIIZILJ:Z return v0 .end method """ class SourceRewriteTest(unittest.TestCase): def test_the_hook_goes_in_front_of_the_original_body(self): patched = media_patch.rewrite_source(SOURCE) head, tail = patched.split(".method public final LIZIZ(", 1) self.assertIn(" .registers 11\n\n" + media_patch.PICK_HOOK, tail) self.assertTrue(tail.index(media_patch.PICK_HOOK) < tail.index("iget-object v1, p0, LX/0qvH;->LJ")) self.assertEqual(SOURCE.split(".method public final LIZIZ(", 1)[0], head) def test_it_hands_over_to_the_plain_pick_and_returns(self): hook = media_patch.PICK_HOOK self.assertIn("invoke-virtual {p0, p1, v0}, LX/0qvH;->LIZ(%sZ)V" % AWEME, hook) self.assertLess(hook.index("LIZ(" + AWEME), hook.index("return-void")) self.assertEqual(1, sum(1 for line in hook.splitlines() if line.strip() == ":lihont_normal")) self.assertEqual(1, hook.count("if-eqz v0, :lihont_normal")) # only v0 is touched, which is a local in a method with eleven registers for register in ("v1", "v2", "v9", "v10"): self.assertNotIn(register, hook) def test_it_refuses_a_second_pass(self): with self.assertRaises(ValueError): media_patch.rewrite_source(media_patch.rewrite_source(SOURCE)) def test_it_refuses_a_class_without_the_methods(self): with self.assertRaises(ValueError): media_patch.rewrite_source(".class public Lx;\n.super Ljava/lang/Object;\n") with self.assertRaises(ValueError): media_patch.rewrite_source(SOURCE.replace(".method public final LIZ(" + AWEME + "Z)V", ".method public final LIZ(I)V")) @unittest.skipUnless(REAL_SOURCE.exists(), "the disassembled X.0qvH is not in work/") def test_the_real_class_is_rewritten(self): text = REAL_SOURCE.read_text(encoding="utf-8").replace(chr(13) + chr(10), chr(10)) patched = media_patch.rewrite_source(text) self.assertEqual(1, patched.count(media_patch.SWITCH)) self.assertEqual(text.count(".method "), patched.count(".method "), "no method added or lost") class ContextRewriteTest(unittest.TestCase): def test_only_the_logo_getter_changes(self): patched = media_patch.rewrite_context(CONTEXT) self.assertIn(media_patch.LOGO, patched) self.assertIn("iget-boolean v0, p0, LX/0qwE;->LJIJ:Z\n\n invoke-static {v0}, " + media_patch.LOGO, patched) # the other flags, including the AI-content ones, keep TikTok's own answers self.assertIn("iget-boolean v0, p0, LX/0qwE;->LJIJI:Z\n\n return v0", patched) self.assertIn("iget-boolean v0, p0, LX/0qwE;->LJIIZILJ:Z\n\n return v0", patched) self.assertEqual(CONTEXT.count(".method "), patched.count(".method ")) def test_registers_cover_what_the_body_uses(self): body = media_patch.LOGO_BODY self.assertIn(".registers 2", body) for register in ("v1", "v2"): self.assertNotIn(register, body) def test_it_refuses_a_second_pass(self): with self.assertRaises(ValueError): media_patch.rewrite_context(media_patch.rewrite_context(CONTEXT)) def test_it_refuses_a_getter_that_is_not_the_plain_one(self): changed = CONTEXT.replace(" return v0\n.end method\n\n.method public final LJIIJ()Z", " invoke-static {v0}, Lx/Y;->z(Z)V\n return v0\n.end method\n\n.method public final LJIIJ()Z") with self.assertRaises(ValueError): media_patch.rewrite_context(changed) with self.assertRaises(ValueError): media_patch.rewrite_context(".class public Lx;\n.super Ljava/lang/Object;\n") @unittest.skipUnless(REAL_CONTEXT.exists(), "the disassembled X.0qwE is not in work/") def test_the_real_class_is_rewritten(self): text = REAL_CONTEXT.read_text(encoding="utf-8").replace(chr(13) + chr(10), chr(10)) patched = media_patch.rewrite_context(text) self.assertEqual(1, patched.count(media_patch.LOGO)) self.assertEqual(text.count(".method "), patched.count(".method "), "no method added or lost") SAVE = """\ .class public final LX/0qv6; .super LX/0qvW; .method public final LJ(Landroid/content/Context;LX/0qvX;)V .registers 18 .prologue move-object/from16 v2, p2 return-void .end method .method public final LJIL(LX/0bqR;)V .registers 3 return-void .end method """ PHOTO = """\ .class public final Lcom/ss/android/ugc/aweme/share/improve/handler/AwemePhotoDownloadShareAbilityHandler; .super Lcom/ss/android/ugc/aweme/share/handler/TTShareDefaultHandler; .method public final LJIILIIL(Landroid/content/Context;Lcom/ss/android/ugc/aweme/share/base/model/BaseSharePackage;LX/0qrm;Landroid/view/View;LX/01h3;)Ljava/lang/Object; .registers 7 .annotation system Ldalvik/annotation/Signature; .end annotation invoke-super/range {p0 .. p5}, Lx/Y;->LJIILIIL()V move-result-object v0 return-object v0 .end method """ class GateRewriteTest(unittest.TestCase): def test_the_save_gate_comes_first_and_can_end_the_method(self): patched = media_patch.rewrite_save(SAVE) self.assertIn(" .registers 18\n\n" + media_patch.SAVE_GATE, patched) self.assertLess(patched.index(media_patch.SAVE_GATE), patched.index("move-object/from16 v2, p2")) gate = media_patch.SAVE_GATE # the arguments are passed as a range, so registers above v15 (p1, p2 here) are fine self.assertIn("invoke-static/range {p0 .. p2}, " + media_patch.GATE, gate) self.assertEqual(1, gate.count("return-void")) self.assertEqual(1, sum(1 for line in gate.splitlines() if line.strip() == ":lihont_go")) self.assertEqual(1, gate.count("if-eqz v0, :lihont_go")) for register in ("v1", "v2", "v3"): self.assertNotIn(register, gate) self.assertEqual(SAVE.count(".method "), patched.count(".method ")) def test_the_photo_gate_completes_a_taken_over_click_with_true(self): patched = media_patch.rewrite_photo(PHOTO) gate = media_patch.PHOTO_GATE self.assertIn(" .registers 7\n\n" + gate, patched) self.assertIn("sget-object v0, Ljava/lang/Boolean;->TRUE:Ljava/lang/Boolean;\n\n return-object v0", gate) self.assertLess(patched.index(gate), patched.index("invoke-super/range")) def test_second_passes_and_foreign_classes_are_refused(self): with self.assertRaises(ValueError): media_patch.rewrite_save(media_patch.rewrite_save(SAVE)) with self.assertRaises(ValueError): media_patch.rewrite_photo(media_patch.rewrite_photo(PHOTO)) with self.assertRaises(ValueError): media_patch.rewrite_save(".class public Lx;\n.super Ljava/lang/Object;\n") with self.assertRaises(ValueError): media_patch.rewrite_photo(PHOTO.replace("LJIILIIL(", "LJIILIIX(")) def test_the_gates_are_aimed_at_the_menu_class(self): self.assertEqual({"LX/0qv6;", "LX/0qvL;", "LX/0qvu;"}, set(media_patch.SAVE_OWNERS)) self.assertIn("lihont/two/media/DownloadMenu;->gate(", media_patch.GATE) self.assertIn("lihont/two/media/DownloadMenu;->gatePhoto(", media_patch.GATE_PHOTO) REAL_SAVE = ROOT / "work" / "dl" / "smali" / "X" / "0qv6.smali" @unittest.skipUnless(REAL_SAVE.exists(), "the disassembled X.0qv6 is not in work/") def test_the_real_save_action_is_rewritten(self): text = self.REAL_SAVE.read_text(encoding="utf-8").replace(chr(13) + chr(10), chr(10)) patched = media_patch.rewrite_save(text) self.assertEqual(1, patched.count(media_patch.GATE)) self.assertEqual(text.count(".method "), patched.count(".method ")) class RuntimeTest(unittest.TestCase): def test_runtime_targets_are_what_the_rewrites_call(self): everything = (media_patch.rewrite_source(SOURCE) + media_patch.rewrite_context(CONTEXT) + media_patch.rewrite_save(SAVE) + media_patch.rewrite_photo(PHOTO)) for target in media_patch.runtime_targets(): self.assertIn(target, everything) if __name__ == "__main__": unittest.main()