HITCON CTF 2017

sakura.py

from z3 import *


f = open("source.c", "rb")

lines = f.read().split("\n")


input_byte = []

s = Solver()

for i in xrange(20):
  for j in xrange(20):
    name = "x_%d_%d" % (i,j)
    x = Int(name)
    input_byte.append(x)
    s.add(x > 0, x <= 9)

assert len(input_byte) == 400

import re
table = [-1]*1200
def build_table(lines):
  prog = re.compile('^  v[0-9]{3,4}\ = [0-9]')

  for l in lines:
    if prog.match(l):
      num = int(l.split(" ")[2][1:]) - 903
      val = int(l.split(" ")[4].split(";")[0])
      table[num] = val

    pass
build_table(lines)


DEBUG = True 
def debug(s):

  if DEBUG:
    print "\tDEBUG: %s" % s

def ppmodel(m):

  for i, k in enumerate(input_byte):
    print "%s = %s" % (k,m[k])


for i, l in enumerate(lines):
  try:
    if "while" in l:
      start_table_index = int(lines[i-2].split("&")[1].split(")")[0][1:]) - 903
      debug("start table index is %d" % start_table_index)
      end_table_index = (int(lines[i-1].split("plus_")[1][:2]) // 4) + start_table_index
      debug("end table index is %d" % end_table_index)
      
      target_val = lines[i+15] 
      if not "if ( accumulator" in target_val:
        target_val = lines[i+13]
      assert "if ( accumulator" in target_val

      target_val = int(target_val.split(" ")[-2])
      debug("target val is %d\n" % target_val)
   
      indices = []
      for i in xrange(start_table_index, end_table_index, 2):
        index = table[i]*20 + table[i+1]  
        #debug("index is %d" % index)
        indices.append(index)

      s.add(Sum([input_byte[x] for x in indices]) == target_val)
      
      for i in indices:
        for j in indices:
          if id(input_byte[i]) != id(input_byte[j]):
            s.add(input_byte[i] != input_byte[j])
      #print s.check()
      #ppmodel(s.model())
  except:
    print "Exception line %d" % i
    break



def generate_output(m):
  l = []
  i = 0
  s = ""
  for k in input_byte:
    if i == 20:
      i=0
      l.append(s)
      s=""
    s += "%s" % m[k]
    i += 1


  i=0
  l.append(s)
  s=""
  return l
print s.check()
payload = generate_output(s.model())


print "".join(payload)

f.close()