import ida_bytes
import ida_segment
import ida_typeinf
import idautils
import ida_xref
import idaapi
from collections import namedtuple

PTR_SIZE = 8 if idaapi.inf_is_64bit() else 4


TypeInfo = namedtuple('SizeInfo', ['get_data', 'typename', 'badaddr'])

SIZE_TO_TYPE_INFO_MAP = {
    1: (TypeInfo(ida_bytes.get_byte,  "byte",  0xFF)),
    2: (TypeInfo(ida_bytes.get_word,  "word",  0xFFFF)),
    4: (TypeInfo(ida_bytes.get_dword, "dword", 0xFFFFFFFF)),
    8: (TypeInfo(ida_bytes.get_qword, "qword", 0xFFFFFFFFFFFFFFFF))
}


def apply_const_type(ea: int) -> bool:
    flags = ida_bytes.get_flags(ea)
    if not ida_bytes.is_data(flags):
        return False

    item_size = ida_bytes.get_item_size(ea)
    type_info = SIZE_TO_TYPE_INFO_MAP.get(item_size)
    if type_info is None:
        return False

    if type_info.get_data(ea) == type_info.badaddr:
        return False

    typeinfo_str = "const " + type_info.typename
    tif = ida_typeinf.tinfo_t(typeinfo_str)

    ida_bytes.clr_op_type(ea, 0)
    return ida_typeinf.apply_tinfo(ea, tif, 1)


def is_addr_readonly(target_addr: int) -> bool:
    flags = ida_bytes.get_full_flags(target_addr)
    return ida_bytes.has_xref(flags) and \
           all(ref.type == ida_xref.dr_R for ref in idautils.XrefsTo(target_addr))


def apply_const_type_if_readonly(target_addr: int) -> bool:
    return apply_const_type(target_addr) if is_addr_readonly(target_addr) else False


def data_is_offset(ea:int) -> bool: 
    item_size = ida_bytes.get_item_size(ea)
    if item_size != PTR_SIZE:
        return False
    
    type_info = SIZE_TO_TYPE_INFO_MAP.get(item_size)
    if type_info is None:
        return False
    
    value = type_info.get_data(ea)
    return  (idaapi.getseg(value)) is not None 


def main():
    print("Casting readonly memory values to consts ")
    print("---------------------------------")
    
    segment_name = ".data"
    
    seg = ida_segment.get_segm_by_name(segment_name)
    if not seg:
        print("No {segment_name} segment")
        return

    ea = seg.start_ea
    fixed_consts_num = 0

    while ea < seg.end_ea:
        ea_size = ida_bytes.get_item_size(ea)
        if ea_size <= 0:
            ea += 1
            continue

        if not data_is_offset(ea) and apply_const_type_if_readonly(ea):
            fixed_consts_num += 1

        ea += ea_size

    print("---------------------------------")
    print(f"Done. fixed = {fixed_consts_num} consts")
    print("Decompile again to view results: F5")


if __name__ == "__main__":
    main()