#!/usr/bin/env python3
"""
Comprehensive GSP firmware analysis using call targets and pattern matching.
"""

import struct
import json
import sys
import os

FWIMAGE_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "gsp-570.144.fwimage.bin")
CALL_TARGETS_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "all_call_targets.txt")
OUTPUT_DIR = os.path.dirname(os.path.abspath(__file__))
BASE_ADDR = 0x04000000

EXEC_SECTIONS = [
    (0x04005000, 0x04006413),
    (0x04007000, 0x04008413),
    (0x04009000, 0x04009073),
    (0x0400a000, 0x0400cfff),
    (0x0400e000, 0x04010fff),
    (0x04018000, 0x0401cfff),
    (0x04c58000, 0x05b21fff),
]

RPC_TYPES = {
    0x01: "ALLOC", 0x02: "FREE", 0x03: "MAP", 0x04: "UNMAP",
    0x05: "PIN", 0x06: "UNPIN", 0x07: "SCHEDULE", 0x08: "YIELD",
    0x09: "NOTIFY", 0x0A: "BAR_TRAP", 0x0B: "MIGRATE", 0xFF: "DEBUG",
}

def read_binary(path):
    with open(path, 'rb') as f:
        return f.read()

def is_in_exec(va):
    for start, end in EXEC_SECTIONS:
        if start <= va <= end:
            return True
    return False

def load_call_targets(path):
    targets = {}
    with open(path, 'r') as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith('#'):
                continue
            parts = line.split()
            if len(parts) >= 1:
                addr = int(parts[0], 16)
                count = int(parts[1]) if len(parts) > 1 else 0
                targets[addr] = count
    return targets

def decode_instruction(data, offset):
    if offset + 4 > len(data):
        return None, 0
    first16 = struct.unpack('<H', data[offset:offset+2])[0]
    if (first16 & 0x3) != 0x3:
        return decode_compressed(first16, offset), 2
    if offset + 4 <= len(data):
        inst = struct.unpack('<I', data[offset:offset+4])[0]
        return decode_32bit(inst, offset), 4
    return None, 0

def decode_compressed(inst, offset):
    if (inst & 0xF003) == 0x8002:
        rd = (inst >> 7) & 0x7; rs2 = (inst >> 10) & 0x7
        return {'mnemonic': 'c.mv', 'rd': rd + 8, 'rs2': rs2 + 8, 'size': 2}
    if (inst & 0xE003) == 0x6001:
        rd = (inst >> 7) & 0x7
        imm = ((inst >> 2) & 0x3F) | ((inst >> 12) & 0x1) << 6
        if imm & 0x40: imm -= 0x80
        return {'mnemonic': 'c.addi', 'rd': rd + 8, 'imm': imm, 'size': 2}
    if (inst & 0xF07F) == 0x8067:
        rs1 = (inst >> 7) & 0x1F
        return {'mnemonic': 'c.jalr', 'rs1': rs1, 'size': 2}
    return {'mnemonic': 'unknown', 'raw': inst, 'size': 2}

def decode_32bit(inst, offset):
    opcode = inst & 0x7F
    funct3 = (inst >> 12) & 0x7
    rd = (inst >> 7) & 0x1F
    rs1 = (inst >> 15) & 0x1F
    rs2 = (inst >> 20) & 0x1F
    if opcode == 0x13 and funct3 == 0:
        imm = (inst >> 20) & 0xFFF
        if imm >= 0x800: imm -= 0x1000
        return {'mnemonic': 'addi', 'rd': rd, 'rs1': rs1, 'imm': imm, 'size': 4}
    if opcode == 0x23 and funct3 == 3:
        offset_val = ((inst >> 7) & 0x1F) | ((inst >> 25) & 0x7F) << 5
        return {'mnemonic': 'sd', 'rs1': rs1, 'rs2': rs2, 'offset': offset_val, 'size': 4}
    if opcode == 0x03 and funct3 == 3:
        offset_val = (inst >> 20) & 0xFFF
        if offset_val >= 0x800: offset_val -= 0x1000
        return {'mnemonic': 'ld', 'rd': rd, 'rs1': rs1, 'offset': offset_val, 'size': 4}
    if opcode == 0x6F:
        imm = ((inst >> 31) & 0x1) << 20 | ((inst >> 12) & 0xFF) << 12 | ((inst >> 20) & 0x1) << 11 | ((inst >> 21) & 0x3FF) << 1
        if imm & 0x100000: imm -= 0x200000
        return {'mnemonic': 'jal', 'rd': rd, 'imm': imm, 'size': 4}
    if opcode == 0x67 and funct3 == 0:
        imm = (inst >> 20) & 0xFFF
        if imm >= 0x800: imm -= 0x1000
        return {'mnemonic': 'jalr', 'rd': rd, 'rs1': rs1, 'imm': imm, 'size': 4}
    if opcode == 0x63:
        imm = ((inst >> 31) & 0x1) << 12 | ((inst >> 7) & 0x1) << 11 | ((inst >> 25) & 0x3F) << 5 | ((inst >> 8) & 0xF) << 1
        if imm & 0x1000: imm -= 0x2000
        return {'mnemonic': 'branch', 'funct3': funct3, 'rs1': rs1, 'rs2': rs2, 'imm': imm, 'size': 4}
    if opcode == 0x17:
        imm = (inst >> 12) & 0xFFFFF
        return {'mnemonic': 'auipc', 'rd': rd, 'imm': imm, 'size': 4}
    if opcode == 0x37:
        imm = (inst >> 12) & 0xFFFFF
        return {'mnemonic': 'lui', 'rd': rd, 'imm': imm, 'size': 4}
    return {'mnemonic': 'unknown', 'raw': inst, 'size': 4}

def scan_function(data, start_offset, max_bytes=2000, call_count=0):
    va = BASE_ADDR + start_offset
    results = {
        'address': va,
        'address_str': f"0x{va:08x}",
        'call_count': call_count,
        'stack_size': None,
        'saves_ra': False,
        'ra_save_reg': None,
        'ra_restore_reg': None,
        'saves_s_regs': [],
        'calls': [],
        'rpc_type_refs': [],
        'has_ret': False,
        'ret_offsets': [],
    }
    i = 0
    in_prologue = True
    while i < max_bytes:
        offset = start_offset + i
        if offset >= len(data): break
        inst_info, size = decode_instruction(data, offset)
        if inst_info is None:
            i += 2; continue
        mnemonic = inst_info.get('mnemonic')
        if in_prologue and mnemonic == 'addi' and inst_info['rd'] == 2 and inst_info['rs1'] == 2 and inst_info['imm'] < -16:
            results['stack_size'] = -inst_info['imm']
            in_prologue = False
        if mnemonic == 'sd' and inst_info['rs1'] == 2 and 9 <= inst_info['rs2'] <= 18:
            results['saves_s_regs'].append(inst_info['rs2'])
        if mnemonic == 'addi' and inst_info['rs1'] == 1 and inst_info['imm'] == 0 and 9 <= inst_info['rd'] <= 18:
            if results['ra_save_reg'] is None:
                results['ra_save_reg'] = inst_info['rd']; results['saves_ra'] = True
        if mnemonic == 'c.mv' and inst_info['rs2'] == 1 and 8 <= inst_info['rd'] <= 15:
            if results['ra_save_reg'] is None:
                results['ra_save_reg'] = inst_info['rd']; results['saves_ra'] = True
        if mnemonic == 'addi' and inst_info['rd'] == 1 and inst_info['imm'] == 0 and 9 <= inst_info['rs1'] <= 18:
            results['ra_restore_reg'] = inst_info['rs1']
        if mnemonic == 'c.mv' and inst_info['rd'] == 1 and 8 <= inst_info['rs2'] <= 15:
            results['ra_restore_reg'] = inst_info['rs2']
        if mnemonic in ('jal', 'jalr'):
            if mnemonic == 'jal':
                results['calls'].append(va + i + inst_info['imm'])
        if mnemonic in ('addi', 'andi', 'ori') and 'imm' in inst_info:
            imm = inst_info['imm']
            if imm in RPC_TYPES:
                results['rpc_type_refs'].append((i, imm, RPC_TYPES[imm]))
        if mnemonic in ('jalr', 'c.jalr') and inst_info.get('rd', 1) == 0 and inst_info.get('imm', 0) == 0 and inst_info.get('rs1', 1) == 1:
            results['has_ret'] = True; results['ret_offsets'].append(i)
        i += size
    return results

def scan_for_strings_in_range(data, start_offset, max_len=500):
    strings = []
    i = start_offset; end = min(start_offset + max_len, len(data))
    while i < end:
        if 32 <= data[i] <= 126:
            start = i
            while i < end and 32 <= data[i] <= 126: i += 1
            if i - start >= 4:
                s = data[start:i].decode('ascii', errors='ignore')
                strings.append((BASE_ADDR + start, s))
        else: i += 1
    return strings

def convert_func(f):
    return {
        'address': f['address_str'],
        'call_count': f.get('call_count', 0),
        'stack_size': f['stack_size'],
        'saves_ra': f['saves_ra'],
        'ra_save_reg': f['ra_save_reg'],
        'ra_restore_reg': f['ra_restore_reg'],
        'saves_s_regs': f['saves_s_regs'],
        'calls': [f"0x{c:08x}" for c in f['calls']],
        'rpc_type_refs': [(RPC_TYPES[r[1]], r[0]) for r in f['rpc_type_refs']],
        'has_ret': f['has_ret'],
    }

def main():
    print("Reading firmware...")
    data = read_binary(FWIMAGE_PATH)
    print(f"Binary size: {len(data)} bytes")
    
    print("Loading call targets...")
    call_targets = load_call_targets(CALL_TARGETS_FILE)
    print(f"Loaded {len(call_targets)} call targets")
    
    sorted_targets = sorted(call_targets.items(), key=lambda x: x[1], reverse=True)
    
    print("\nScanning top 500 call targets...")
    functions = []
    for addr, count in sorted_targets[:500]:
        if not is_in_exec(addr): continue
        offset = addr - BASE_ADDR
        if offset < 0 or offset >= len(data): continue
        func_info = scan_function(data, offset, max_bytes=1000, call_count=count)
        functions.append(func_info)
    
    print(f"Analyzed {len(functions)} functions")
    
    rpc_handlers = [f for f in functions if f['rpc_type_refs']]
    other_funcs = [f for f in functions if not f['rpc_type_refs']]
    
    print(f"\nRPC handler candidates: {len(rpc_handlers)}")
    print(f"Other functions: {len(other_funcs)}")
    
    print("\n=== Top 20 RPC Handler Candidates ===")
    for f in rpc_handlers[:20]:
        rpc_refs = [(RPC_TYPES[r[1]], r[0]) for r in f['rpc_type_refs']]
        print(f"  {f['address_str']}: calls={f['call_count']}, stack={f['stack_size']}, rpc_refs={rpc_refs[:10]}")
    
    # Save results
    results = {
        'total_functions_analyzed': len(functions),
        'rpc_handlers': [convert_func(f) for f in rpc_handlers[:200]],
        'other_functions': [convert_func(f) for f in other_funcs[:100]],
    }
    with open(f"{OUTPUT_DIR}/analysis_results.json", "w") as f:
        json.dump(results, f, indent=2)
    print(f"\nResults saved to {OUTPUT_DIR}/analysis_results.json")
    
    # Detailed scans
    print("\n=== Entry Point (0x04009000) ===")
    print(json.dumps(convert_func(scan_function(data, 0x9000, max_bytes=500)), indent=2))
    
    print("\n=== RPC_BAR_TRAP Handler (0x04cb8189) ===")
    print(json.dumps(convert_func(scan_function(data, 0xcb8189, max_bytes=500)), indent=2))
    
    print("\n=== Top 3 Functions Detail ===")
    for f in functions[:3]:
        print(f"\n--- {f['address_str']} (calls={f['call_count']}) ---")
        print(json.dumps(convert_func(f), indent=2))

if __name__ == "__main__":
    main()