import logging

FORMAT = "[%(levelname)s] %(name)s %(filename)s:%(lineno)d: %(message)s"
logging.basicConfig(format=FORMAT)
logging.getLogger().setLevel(logging.DEBUG)

import json
import androscalpel as asc
import zipfile as z
from androscalpel import *

# APK_NAME = "test.apk"
APK_NAME = __file__.removesuffix("test.py") + "/apk_frauder/app-release.apk"
DEX_NAME = "classes.dex"

with z.ZipFile(APK_NAME) as zipf:
    with zipf.open(DEX_NAME, "r") as dex:
        dex = dex.read()

apk = asc.Apk()
apk.add_dex_file(dex)

clazz_id = IdType("Lcom/example/testapplication/ui/home/HomeViewModel;")
proto_id = IdMethodType(IdType("Ljava/lang/String;"), [])
method_id = IdMethod("text_gen", proto_id, clazz_id)

clazz = apk.classes[clazz_id]
method = clazz.virtual_methods[method_id]
code = method.code

logging.getLogger().setLevel(logging.ERROR)

print(f"[+] Code of {method_id} ")
for i in code.insns:
    print(f"    {i}")
print("[+] Modify code")
new_insns = []
for i in code.insns:
    if isinstance(i, asc.ins.ConstString):
        if i.lit == "Hello":
            i = asc.ins.ConstString(i.reg, DexString("Degemer Mat"))
        elif i.lit == "Bye":
            i = asc.ins.ConstString(i.reg, DexString("Kenavo"))
    new_insns.append(i)

# This need improving!
code = asc.Code(code.registers_size, code.ins_size, code.outs_size, new_insns)
apk.set_method_code(method_id, code)
# apk.set_method_code(method.descriptor, code)


clazz = apk.classes[clazz_id]
method = clazz.virtual_methods[method_id]
code = method.code
print(f"[+] New code of {method_id} ")
for i in code.insns:
    print(f"    {i}")

# Strip class for debugging
# classes = list(
#    filter(
#        lambda x: x
#        not in [
#            IdType("Lcom/example/testapplication/ui/home/HomeViewModel;"),
#            IdType("Landroidx/navigation/NavDeepLink$Builder;"),
#            IdType("Landroidx/constraintlayout/core/widgets/ConstraintWidget$1;"),
#        ],
#        apk.classes.keys(),
#    )
# )
# for cls in classes:
#    apk.remove_class(cls)

print("[+] Recompile")

dex_raw = apk.gen_raw_dex()
assert len(dex_raw) == 1
with open(DEX_NAME, "wb") as file:
    file.write(dex_raw[0])

print("[+] Load new dex")
with open(DEX_NAME, "rb") as file:
    dex = file.read()
new_apk = asc.Apk()
new_apk.add_dex_file(dex)

clazz = new_apk.classes[clazz_id]
method = clazz.virtual_methods[method_id]
code = method.code

print(f"[+] Code of {method_id} in new apk")
for i in code.insns:
    print(f"    {i}")