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 disassembly is partially gibberish
# - the program calls an enc procedure to decrypt itself, then calls a
#   procedure at 0x2400 (presumably inside the decrypted parts)
# - trace execution from 0x2400 onwards to see what the program is doing

# 4438 <main>
# 4438:  3e40 2045      mov	#0x4520, r14
# 443c:  0f4e           mov	r14, r15
# 443e:  3e40 f800      mov	#0xf8, r14
# 4442:  3f40 0024      mov	#0x2400, r15
# 4446:  b012 8644      call	#0x4486 <enc>
# 444a:  b012 0024      call	#0x2400
# 444e:  0f43           clr	r15


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


def trace(password):
    dbg.sendlinethen(PROMPT, b'break *0x444a')
    dbg.sendlinethen(PROMPT, b'continue')
    dbg.sendlinethen(PROMPT, b'si')  # jump over breakpoint instruction
    emu.sendline(password)

    output = None
    while output != 'No registers.':
        output = gdb_output(dbg.sendlinethen(PROMPT, b'x/i $pc'))[0]
        print(output)
        dbg.sendlinethen(PROMPT, b'si')

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

# trace(b'A'*32)

# the trace shows loads of instructions executed to print a password
# prompt, followed by reading and checking the password:

# => 0x2444:    add    #6,    r1    ;#0x0006
# => 0x2448:    cmp    #19058,    -36(r4)    ;#0x4a72, 0xffdc(r4)
# => 0x244e:    jnz    $+12         ;abs 0x245a
# => 0x245a:    add    #32,    r1    ;#0x0020
# => 0x245e:    pop    r4
# => 0x2460:    pop    r11
# => 0x2462:    ret
# => 0x444e:    clr    r15
# => 0x4450:    bis    #240,    r2    ;#0x00f0
# No registers.

# let's check what's in the memory at 0x2448


def check_memory(payload):
    # break after decryption for 0x2444 to be a valid address
    dbg.sendlinethen(PROMPT, b'break *0x444a')
    dbg.sendlinethen(PROMPT, b'continue')
    dbg.sendlinethen(PROMPT, b'break *0x2444')
    dbg.sendline(b'continue')

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

    dbg.recvuntil(PROMPT)

    print(gdb_output(dbg.sendlinethen(PROMPT, b'x/4xb $r4-36'))[0])


# check_memory(b'ABCD')

# 0x43da:    0x41    0x42    0x43    0x44

# surprisingly enough, it's our input, therefore it just needs to be
# matched with the constant it's compared to

def exploit():
    payload = p16(0x4a72)

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

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


exploit()