from pathlib import Path


def u16(data, off):
    return data[off] | (data[off + 1] << 8)


def u32(data, off):
    return u16(data, off) | (u16(data, off + 2) << 16)


def selector_to_segment_index(selector):
    return ((selector - 0x1000) // 8) + 1


def get_segment_info(data, selector):
    ne_off = u32(data, 0x3C)
    seg_table_off = ne_off + u16(data, ne_off + 0x22)
    align_shift = u16(data, ne_off + 0x32)
    seg_index = selector_to_segment_index(selector)
    seg_entry_off = seg_table_off + (seg_index - 1) * 8
    sector = u16(data, seg_entry_off)
    seg_len = u16(data, seg_entry_off + 2)
    if seg_len == 0:
        seg_len = 0x10000
    flags = u16(data, seg_entry_off + 4)
    seg_file_off = sector << align_shift
    reloc_off = seg_file_off + seg_len
    reloc_count = u16(data, reloc_off)
    return {
        'selector': selector,
        'segment_index': seg_index,
        'segment_file_off': seg_file_off,
        'segment_len': seg_len,
        'flags': flags,
        'reloc_off': reloc_off,
        'reloc_count': reloc_count,
    }


def read_relocs(data, seg_info):
    relocs = []
    pos = seg_info['reloc_off'] + 2
    for i in range(seg_info['reloc_count']):
        relocs.append({
            'index': i,
            'record_off': pos,
            'src_type': data[pos],
            'flags': data[pos + 1],
            'src_off': u16(data, pos + 2),
            'b4': data[pos + 4],
            'b5': data[pos + 5],
            'target_off': u16(data, pos + 6),
            'raw': data[pos:pos + 8],
        })
        pos += 8
    return relocs


def dump_site(data, seg_info, site, size=16):
    raw = seg_info['segment_file_off'] + site
    print(f'site {seg_info["selector"]:04X}:{site:04X} -> file 0x{raw:X}: {data[raw:raw+size].hex(" ")}')


def main():
    path = Path(r'd:\Ghidra\Crusader\REGRET.EXE')
    data = path.read_bytes()

    for selector in (0x1148, 0x13F8):
        info = get_segment_info(data, selector)
        print(f'selector 0x{selector:04X} -> segment index {info["segment_index"]}, file 0x{info["segment_file_off"]:X}, len 0x{info["segment_len"]:X}, reloc@0x{info["reloc_off"]:X}, count={info["reloc_count"]}')
        relocs = read_relocs(data, info)
        interesting = []
        if selector == 0x1148:
            interesting = [0x3678, 0x36C0, 0x3702, 0x3743]
        else:
            interesting = [0x10FA, 0x2040, 0x209F, 0x20B9]
        for site in interesting:
            dump_site(data, info, site, 24)
            local_hits = [r for r in relocs if site <= r['src_off'] <= site + 0x80]
            for r in local_hits[:20]:
                print(f'  reloc idx={r["index"]} rec@0x{r["record_off"]:X} src=0x{r["src_off"]:04X} bytes={r["raw"].hex(" ")}')
        print('')


if __name__ == '__main__':
    main()