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 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):
            print(f"{i:>03}    {inst}")
            print(
                f"        val: {GREEN}{dyn_load_apk.classes[field.class_].static_fields[field].value}{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]
#
# new_insns = (
#    code.insns[:42]
#    + [
#        ins.NewInstance(1, IdType("Lcom/example/ut_dyn_load/SmsReceiver;")),
#        ins.InvokeVirtual(m_id, [1, 8, 1]),
#    ]
#    + code.insns[62:]
# )
#
# print(f"[+] New code ")
# for i, inst in enumerate(new_insns):
#    if i >= 42 and i < 44:
#        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}")
#            print(
#                f"        val: {GREEN}{dyn_load_apk.classes[field.class_].static_fields[field].value}{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, code)

NB_APK = 8665

list_cls = list(dyn_load_apk.classes.keys())
list_cls.sort()
print(f"[+] NB classes: {len(list_cls)}")
for cls in list_cls[NB_APK:]:
    dyn_load_apk.remove_class(cls)

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


MAX_REQ = 8


def cmp(a, b, req=0):
    if req > MAX_REQ:
        return
    if type(a) == dict:
        cmp_dict(a, b, req)
    elif type(a) == list:
        cmp_list(a, b, req)
    else:
        cmp_other(a, b, req)


def nice_bool(b) -> str:
    if b:
        return "\033[32mTrue\033[0m"
    else:
        return "\033[31mFalse\033[0m"


def cmp_other(a, b, req=0):
    ident = "  " * req
    for f in dir(a):
        if getattr(getattr(a, f), "__call__", None) is None and (
            len(f) < 2 or f[:2] != "__"
        ):
            eq = getattr(a, f) == getattr(b, f)
            if not eq:
                print(f"{f'{ident}{f}: ':<150}{nice_bool(eq)}")
                if "descriptor" in dir(a):
                    global last_id
                    last_id = a.descriptor
                cmp(getattr(a, f), getattr(b, f), req + 1)


def cmp_dict(a, b, req=0):
    ident = "  " * req
    keys_a = set(a.keys())
    keys_b = set(b.keys())
    if keys_a != keys_b:
        print(f"{ident}a.keys() != b.keys()")
    tot = 0
    nb_failed = 0
    for key in keys_a & keys_b:
        eq = a[key] == b[key]
        tot += 1
        if not eq:
            nb_failed += 1
            print(f"{f'{ident}{str(key)}: ':<150}{nice_bool(eq)}")
            global last_id
            last_id = key
            cmp(a[key], b[key], req + 1)
    print(f"\033[32m{tot-nb_failed}\033[0m + \033[31m{nb_failed}\033[0m = {tot}")


def cmp_list(a, b, req=0):
    ident = "  " * req
    la = len(a)
    lb = len(b)
    if la != lb:
        print(f"{ident}len(a) != len(b)")
    for i in range(min(la, lb)):
        eq = a[i] == b[i]
        if not eq:
            print(f"{f'{ident}{str(i)}: ':<150}{nice_bool(eq)}")
            print(f"{ident}: {str(a[i])} != {str(b[i])}")
            cmp(a[i], b[i], req + 1)


instrumented_apk = Apk()
instrumented_apk.add_dex_file(dex_raw[0])

cmp(instrumented_apk, dyn_load_apk)

MAX_REQ = 5

tys = [
    IdType("Landroidx/compose/foundation/MutatePriority;"),
    IdType(
        "Landroidx/compose/ui/input/pointer/PointerInteropFilter$DispatchToViewState;"
    ),
    IdType("Landroidx/compose/ui/autofill/AutofillType;"),
    IdType("Landroidx/compose/material3/TextFieldType;"),
    IdType("Landroidx/compose/material3/SnackbarResult;"),
    IdType("Landroidx/compose/ui/layout/MeasuringIntrinsics$IntrinsicMinMax;"),
    IdType("Landroidx/compose/foundation/text/Handle;"),
    IdType("Landroidx/compose/foundation/text/selection/SelectionMode;"),
    IdType("Landroidx/compose/material3/SliderComponents;"),
    IdType("Landroidx/lifecycle/Lifecycle$Event;"),
    IdType("Landroidx/compose/ui/focus/FocusStateImpl;"),
    IdType("Landroidx/fragment/app/strictmode/FragmentStrictMode$Flag;"),
    IdType("Landroidx/compose/ui/text/style/ResolvedTextDirection;"),
    IdType("Landroidx/compose/foundation/text/KeyCommand;"),
    IdType("Landroidx/compose/ui/text/AnnotationType;"),
    IdType("Landroidx/compose/foundation/layout/SizeMode;"),
    IdType("Landroidx/compose/ui/semantics/NodeLocationHolder$ComparisonStrategy;"),
    IdType("Landroidx/compose/material3/DrawerValue;"),
    IdType("Landroidx/compose/ui/input/pointer/PointerEventPass;"),
    IdType("Landroidx/compose/foundation/layout/Direction;"),
    IdType("Landroidx/compose/ui/node/LayoutNode$UsageByParent;"),
    IdType("Landroidx/compose/material3/tokens/TypographyKeyTokens;"),
    IdType("Landroidx/compose/animation/core/RepeatMode;"),
    IdType("Landroidx/compose/ui/layout/IntrinsicMinMax;"),
    IdType("Landroidx/compose/ui/text/input/TextInputServiceAndroid$TextInputCommand;"),
    IdType("Landroidx/compose/ui/unit/LayoutDirection;"),
    IdType("Landroidx/collection/SparseArrayCompat;"),
    IdType("Landroidx/compose/material3/TabSlots;"),
    IdType("Landroidx/compose/material3/tokens/ColorSchemeKeyTokens;"),
    IdType("Landroidx/annotation/RestrictTo$Scope;"),
    IdType("Landroidx/profileinstaller/FileSectionType;"),
    IdType("Landroidx/lifecycle/Lifecycle$State;"),
    IdType("Landroidx/compose/animation/EnterExitState;"),
    IdType("Landroidx/compose/animation/core/MutatePriority;"),
    IdType("Landroidx/compose/ui/layout/MeasuringIntrinsics$IntrinsicWidthHeight;"),
    IdType("Landroidx/loader/content/ModernAsyncTask$Status;"),
    IdType("Landroidx/annotation/InspectableProperty$ValueType;"),
    IdType("Landroidx/compose/foundation/layout/LayoutOrientation;"),
    IdType("Landroidx/compose/ui/node/NodeMeasuringIntrinsics$IntrinsicMinMax;"),
    IdType("Landroidx/compose/material3/tokens/ShapeKeyTokens;"),
    IdType("Landroidx/compose/ui/node/LayoutNode$LayoutState;"),
    IdType("Landroidx/compose/runtime/InvalidationResult;"),
    IdType("Landroidx/compose/ui/layout/IntrinsicWidthHeight;"),
    IdType("Landroidx/compose/ui/node/NodeMeasuringIntrinsics$IntrinsicWidthHeight;"),
    IdType("Landroidx/compose/runtime/Recomposer$State;"),
    IdType("Landroidx/compose/foundation/gestures/Orientation;"),
    IdType("Landroidx/compose/material3/ScaffoldLayoutContent;"),
    IdType("Landroidx/compose/foundation/text/selection/HandleReferencePoint;"),
    IdType("Landroidx/compose/material3/SnackbarDuration;"),
    IdType("Landroidx/compose/ui/platform/TextToolbarStatus;"),
    IdType("Landroidx/compose/material3/InputPhase;"),
    IdType("Landroidx/compose/ui/window/SecureFlagPolicy;"),
    IdType("Landroidx/collection/LongSparseArray;"),
    IdType("Landroidx/compose/animation/core/AnimationEndReason;"),
    IdType("Landroidx/compose/foundation/text/HandleState;"),
    IdType("Landroidx/compose/ui/state/ToggleableState;"),
    IdType("Landroidx/compose/ui/text/android/animation/SegmentType;"),
    IdType("Landroidx/compose/ui/platform/actionmodecallback/MenuItemOption;"),
    IdType("Landroidx/compose/foundation/layout/IntrinsicSize;"),
    IdType("Landroidx/compose/ui/platform/actionmodecallback/MenuItemOption;"),
]
instrumented_apk_classes = instrumented_apk.classes
dyn_load_apk_classes = dyn_load_apk.classes
for ty in tys:
    if instrumented_apk_classes[ty].annotations != dyn_load_apk_classes[ty].annotations:
        print(f"-> {str(ty)}")
        # invoke-virtual {0} Landroid/animation/Animator;->end()V != invoke-virtual {0} [Ljava/lang/Object;->clone()Ljava/lang/Object;
        cmp(
            instrumented_apk.classes[ty],
            dyn_load_apk.classes[ty],
        )