import os
import re
from z3 import *

# ================= CONFIGURATION =================
# Path to the folder containing the decompiled .java files
SOURCE_DIR = "./src_files/sources/com/lake/ctf" 

# Initial Entry Point (From the Native Init function)
START_CLASS_HASH = "86e50149658661312a9e0b35558d84f6c6d3da797f552a9657fe0558ca40cdef"
START_METHOD_HASH = "031b4af5197ec30a926f48cf40e11a7dbc470048a21e4003b7a3c07c5dab1baa"

# ================= Z3 SETUP =================
# UPDATED: Flag length is explicitly 55 based on MainActivity
FLAG_LEN = 55
# UPDATED: The app checks exactly 80 links in the chain
MAX_STEPS = 80

solver = Solver()
flag_vars = [BitVec(f"flag_{i}", 8) for i in range(FLAG_LEN)]

# Add printable ASCII constraints
for i in range(FLAG_LEN):
    solver.add(flag_vars[i] >= 32)
    solver.add(flag_vars[i] <= 126)

# ================= HELPER CLASSES =================

class MockString:
    """
    A helper class to allow us to 'eval' the Java equation directly.
    """
    def charAt(self, index):
        if index >= FLAG_LEN:
            # This is now a critical error since we know the exact length
            print(f"[!] Critical: Equation references index {index} which is >= 55.")
            return 0
        
        # We must zero-extend to 32 bits because Java intermediate math is int (32-bit)
        # but our chars are 8-bit.
        return ZeroExt(24, flag_vars[index]) 

# ================= MAIN LOGIC =================

def load_files():
    """Reads all Java files into a dictionary keyed by Class Hash."""
    files_map = {}
    if not os.path.exists(SOURCE_DIR):
        print(f"Error: Directory {SOURCE_DIR} not found.")
        return {}

    print(f"[*] Loading files from {SOURCE_DIR}...")
    for f in os.listdir(SOURCE_DIR):
        if f.endswith(".java") and f.startswith("Check"):
            class_hash = f.replace("Check", "").replace(".java", "")
            with open(os.path.join(SOURCE_DIR, f), "r") as file_obj:
                files_map[class_hash] = file_obj.read()
    print(f"[*] Loaded {len(files_map)} files.")
    return files_map

def extract_logic(class_hash, method_hash, files_map):
    content = files_map.get(class_hash)
    if not content:
        print(f"[!] Critical: Class hash {class_hash} not found on disk.")
        return None, None, None

    # Pattern: static boolean Check[METHOD_HASH](String str) { BODY }
    method_pattern = re.compile(rf"static boolean Check{method_hash}\(String (\w+)\)\s*\{{(.*?)\}}", re.DOTALL)
    match = method_pattern.search(content)
    
    if not match:
        print(f"[!] Critical: Method {method_hash} not found in class {class_hash}.")
        return None, None, None

    param_name = match.group(1)
    body = match.group(2)

    # Extract the 'if' condition
    if_match = re.search(r"if\s*\((.*)\)\s*\{", body)
    if not if_match:
        return None, None, None
    
    equation_str = if_match.group(1)

    # Extract the 'Success' Nop call
    if_block_start = body.find("{")
    if_block_end = body.find("}", if_block_start)
    if_block_content = body[if_block_start:if_block_end]

    nop_match = re.search(r'nop\("([a-f0-9]+)",\s*"([a-f0-9]+)"\)', if_block_content)
    
    if not nop_match:
        return equation_str, None, None
    
    return equation_str, nop_match.group(1), nop_match.group(2)

def solve():
    files_map = load_files()
    if not files_map: 
        return

    curr_class = START_CLASS_HASH
    curr_method = START_METHOD_HASH
    
    step_count = 0
    
    print(f"[*] Starting traversal for exactly {MAX_STEPS} steps...")

    for i in range(MAX_STEPS):
        if not curr_class:
            print("[!] Chain ended prematurely!")
            break
       
        step_count += 1
        
        eq_str, next_class, next_method = extract_logic(curr_class, curr_method, files_map)
        
        if not eq_str:
            break

        try:
            if "==" in eq_str:
                print(f"[*] Adding constraint {eq_str}")
                lhs, rhs = eq_str.split("==")
                ctx = {'str': MockString()} 
                lhs_expr = eval(lhs.strip(), {}, ctx)
                rhs_val = int(eval(rhs.strip(), {}, {}))
                solver.add(lhs_expr == rhs_val)
            else:
                print(f"[!] Warning: Equation format unknown: {eq_str}")
        except Exception:
            print(f"[!] Error parsing equation: {eq_str}")
            break

        curr_class = next_class
        curr_method = next_method

    print(f"[*] Extraction complete. {step_count} constraints added.")
    print("[*] Solving with Z3...")

    if solver.check() == sat:
        model = solver.model()
        result = []
        for i in range(FLAG_LEN):
            val = model[flag_vars[i]]
            if val is not None:
                result.append(chr(val.as_long()))
            else:
                result.append("?")
        
        final_flag = "".join(result)
        print("\n" + "="*50)
        print(f"FLAG: {final_flag}")
        print("="*50 + "\n")
    else:
        print("[!] Unsatisfiable! Constraints might be contradictory.")

if __name__ == "__main__":
    solve()