import logging

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

import json
import zipfile as z
import re
from pathlib import Path

from androscalpel import Apk, IdType, IdMethodType, ins, DexString, IdMethod, Code, utils, Field, IdField  # type: ignore

RED = "\033[38:2:255:0:0m"
GREEN = "\033[38:2:0:255:0m"
ENDC = "\033[0m"

RE_CHECK_DEXFILE = re.compile(r"classes\d*.dex")


def is_dexfile(filename: str) -> bool:
    return bool(RE_CHECK_DEXFILE.fullmatch(filename))


ORIGINAL_APK = Path(__file__).parent / "origin-release.apk"
DYN_LOAD_APK = Path(__file__).parent / "ut-dyn-load-release.apk"
DEX_NAME = "classes.dex"

clazz_id = IdType("Lcom/example/ut_dyn_load/Loader;")
proto_id = IdMethodType(IdType.void(), [IdType("Landroid/content/Context;")])
method_id = IdMethod("load", proto_id, clazz_id)


def load_apk(path_apk: Path) -> Apk:
    print(f"[+] Load bytecode ")
    apk = Apk()
    with z.ZipFile(path_apk) as zipf:
        for file in filter(is_dexfile, zipf.namelist()):
            print(f"[-] {file}")
            with zipf.open(file, "r") as dex_f:
                dex = dex_f.read()
                apk.add_dex_file(dex)
    return apk


print(f"[+] Load bytecode ")
dyn_load_apk = load_apk(DYN_LOAD_APK)

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

# logging.getLogger().setLevel(logging.WARNING)


def is_evasion_method(meth: IdMethod) -> bool:
    return (
        method.class_ == IdType.from_smali("Ldalvik/system/PathClassLoader;")
        or (
            method.class_ == IdType.from_smali("Ljava/lang/Class;")
            and method.name == "newInstance"
        )
        or (
            method.class_ == IdType.from_smali("Ljava/lang/Class;")
            and method.name == "getMethod"
        )
        or method.class_ == IdType.from_smali("Ljava/lang/reflect/Method;")
    )


print(f"[+] Code of {method_id} ")
for i, inst in enumerate(code.insns):
    match inst:
        case ins.InvokeVirtual(args=args, method=method) if is_evasion_method(method):
            print(f"{i:>03}    {RED}{inst}{ENDC}")
        case ins.SGetObject(to=to, field=field):
            field: IdField = field  # type: ignore
            print(f"{i:>03}    {inst}")
            val = dyn_load_apk.classes[field.class_].static_fields[field].value
            print(f"        val: {GREEN}{val}{ENDC}")
        case inst:
            print(f"{i:>03}    {inst}")

malicious_class_id = IdType("Lcom/example/ut_dyn_load/SmsReceiver;")

mal_cls_in_apk = malicious_class_id in dyn_load_apk.classes
if mal_cls_in_apk:
    color = GREEN
else:
    color = RED

print(f"[+] {malicious_class_id} in loaded apk: {color}{mal_cls_in_apk}{ENDC}")

print(f"[+] Load bytecode in assets/classes")
with z.ZipFile(DYN_LOAD_APK) as zipf:
    with zipf.open("assets/classes", "r") as dex_f:
        dex = dex_f.read()
        dyn_load_apk.add_dex_file(dex)

mal_cls_in_apk = malicious_class_id in dyn_load_apk.classes
if mal_cls_in_apk:
    color = GREEN
else:
    color = RED
print(f"[+] {malicious_class_id} in loaded apk: {color}{mal_cls_in_apk}{ENDC}")

print(
    f"[+] -- very simplified steep, see https://developer.android.com/reference/java/lang/Class#getMethod(java.lang.String,%20java.lang.Class%3C?%3E[]) --"
)


name = "a"
args = [
    IdType("Landroid/content/Context;"),
    IdType("Landroid/content/BroadcastReceiver;"),
]
malicious_class = dyn_load_apk.classes[malicious_class_id]

potential_meth = []
for m_id in malicious_class.direct_methods:
    if m_id.name == name and m_id.proto.get_parameters() == args:
        potential_meth.append((m_id, "direct"))
for m_id in malicious_class.virtual_methods:
    if m_id.name == name and m_id.proto.get_parameters() == args:
        potential_meth.append((m_id, "virtual"))

print("[+] Potential methods:")
for m_id, t in potential_meth:
    print(f"    {m_id}({t})")

m_id, t = potential_meth[0]


TRY_INDEX = 80  # try until label_0000007C    Ljava/lang/Exception;: label_0000007D
LABEL_7C_INDEX = 121
code_inserted = [
    ins.NewInstance(1, IdType("Lcom/example/ut_dyn_load/SmsReceiver;")),
    ins.InvokeDirect(
        IdMethod.from_smali("Lcom/example/ut_dyn_load/SmsReceiver;-><init>()V"), [1]
    ),
    ins.InvokeStatic(m_id, [8, 1]),
]

new_insns = code.insns[:TRY_INDEX] + code_inserted + code.insns[LABEL_7C_INDEX:]

print(f"[+] New code ")
for i, inst in enumerate(new_insns):
    if i >= TRY_INDEX and i < TRY_INDEX + len(code_inserted):
        print(f"{i:>03}    {GREEN}{inst}{ENDC}")
        continue
    match inst:
        case ins.InvokeVirtual(args=args, method=method) if is_evasion_method(method):
            print(f"{i:>03}    {RED}{inst}{ENDC}")
        case ins.SGetObject(to=to, field=field):
            print(f"{i:>03}    {inst}")
            val = dyn_load_apk.classes[field.class_].static_fields[field].value
            print(f"        val: {GREEN}{val}{ENDC}")
        case inst:
            print(f"{i:>03}    {inst}")

new_code = Code(code.registers_size, code.ins_size, code.outs_size, new_insns)
dyn_load_apk.set_method_code(method_id, new_code)

print("[+] Recompile")
dex_raw = dyn_load_apk.gen_raw_dex()

print("[+] Repackage")
utils.replace_dex(
    DYN_LOAD_APK,
    DYN_LOAD_APK.parent
    / (DYN_LOAD_APK.name.removesuffix(".apk") + "-instrumented.apk"),
    dex_raw,
    Path(__file__).parent.parent / "my-release-key.jks",
    zipalign=Path.home() / "Android" / "Sdk" / "build-tools" / "34.0.0" / "zipalign",
    apksigner=Path.home() / "Android" / "Sdk" / "build-tools" / "34.0.0" / "apksigner",
)