LihonT/tests/test_media_patch.py

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")


class RuntimeTest(unittest.TestCase):
    def test_runtime_targets_are_what_the_rewrites_call(self):
        both = media_patch.rewrite_source(SOURCE) + media_patch.rewrite_context(CONTEXT)
        for target in media_patch.runtime_targets():
            self.assertIn(target, both)


if __name__ == "__main__":
    unittest.main()