asisctf_odrerre_solve.py

#!/usr/bin/python

import os
import sys
import zlib
import struct
import copy

# ghetto recursive png brute force thingy

__QUIET__ = True

def quiet_print(s):
  if __QUIET__ == False:
    print s

SENSITIVITY = 0x40

# identifies chunks we think are likely
def identify_optimal_next_chunk(zlib_partial,idats_remaining):
  good_tries = []
  for ir in idats_remaining:
    for i in range(0,SENSITIVITY):
      zl_temp = zlib_partial + ir[8:(-4 - i)]
      try:
        zlib.decompress(zl_temp)
        good_tries.append(ir)
        break
      except zlib.error,e:
        if "-5" in str(e):
          good_tries.append(ir)
          break
        elif "-3" in str(e):
          pass
  return good_tries

# brute force based on optimum block length
def bruteforce_partial(headers, idats_known, idats_remaining, end):
  if len(idats_remaining) == 0:
    c = ""
    for x in idats_known:
      c += x[8:-4]
    try:
      zlib.decompress(c)
    except:
      quiet_print("dead end")
      return
    print " [+] no zlib errors, writing to out.png"
    f = open("out.png","wb")
    f.write(headers +  ''.join(idats_known) +  end)
    f.close()
    print " [+] checking with `file out.png`..."
    os.system("file out.png")
    sys.exit(0)
  else:
    zlib_partial = "".join(x[8:-4] for x in idats_known)
    next_list = identify_optimal_next_chunk(zlib_partial,idats_remaining)
    quiet_print("identify_optimal_next_chunk returned %d hits" % len(next_list))
    for next_chunk in next_list:
      new_idats_remaining = copy.copy(idats_remaining)
      new_idats_remaining.remove(next_chunk)
      bruteforce_partial(headers,idats_known + [next_chunk],new_idats_remaining,end)
      del(new_idats_remaining)

# brute force and prune on entire blocks
def bruteforce(headers, idats_known,idats_remaining, end):
  if len(idats_remaining) == 0:
    c = ""
    for x in idats_known:
      c += x[8:-4]
    zlib.decompress(c)
    print " [+] no zlib errors, writing to out.png"
    f = open("out.png","wb")
    f.write(headers +  ''.join(idats_known) +  end)
    f.close()
    print " [+] checking with `file out.png`..."
    os.system("file out.png")
    sys.exit(0)
  else:
    chunk_rebuild = ""
    for i in idats_known:
      chunk_rebuild = i[8:-4]
    failed_decompress = 0
    for i in idats_remaining:
      temp = chunk_rebuild + i[8:-4]
      try:
        zlib.decompress(temp)
      except zlib.error,e:
        if "-5" in str(e):
          quiet_print("optimisation : zlib error -5 at depth %d" % len(idats_known))
          new_idats_known = idats_known + [i]
          new_idats_remaining = copy.copy(idats_remaining)
          new_idats_remaining.remove(i)
          # recursively call self
          bruteforce(headers,new_idats_known,new_idats_remaining,end)
          continue
        elif "-3" in str(e):
          failed_decompress += 1 
          continue
    # all the options have failed, recursively explore options.
    if failed_decompress == len(idats_remaining) and failed_decompress != 1:
      for remain in idats_remaining:
        quiet_print("no optimisation brute forcing a chunk at depth %d, exploring %d remaining options" % (len(idats_known),len(idats_remaining)))
        new_idats_known = idats_known + [remain]
        new_idats_remaining = copy.copy(idats_remaining)
        new_idats_remaining.remove(remain)
        bruteforce(headers,new_idats_known,new_idats_remaining,end)
    elif failed_decompress != 0:
      quiet_print("dead end, returning")
      return
    elif failed_decompress == 0:
      quiet_print("win state, finishing (i.e. there are no)")
      new_idats_known = idats_known + [i]
      new_idats_remaining = copy.copy(idats_remaining)
      new_idats_remaining.remove(i)
      # recursively call self
      bruteforce(headers,new_idats_known,new_idats_remaining,end)
    return # i think this is a failed decompression state
    # print "how did we get here: failed_decompress is %d, known_idats is %d, remaining %d" % (failed_decompress, len(idats_known), len(idats_remaining))

def png_split_chunks_simple(fname):
  chunk_headers = ""
  chunk_idats = []
  chunk_iend = ""
  hdr_count = 0
  idat_count = 0
  f = open(fname,"r")
  hdr = f.read(8)
  print " [+] png header: [%s]" % ':'.join(hex(ord(c)) for c in hdr)
  while True:
    chunk_hdr = f.read(8)
    (d,) = struct.unpack(">i",chunk_hdr[0:4])
    (s,) = struct.unpack("4s",chunk_hdr[4:])
    chunk_data = f.read(d+4)
    # print "   > read chunk %s" % s
    if s == "IDAT":
      idat_count += 1
      chunk_idats.append(chunk_hdr + chunk_data)
    elif s == "IEND":
      print " [+] iend reached : %d headers, %d idats" % (hdr_count, idat_count)
      chunk_iend = chunk_hdr + chunk_data
      break
    else:
      hdr_count += 1
      chunk_headers += chunk_hdr + chunk_data
  f.close()
  return (hdr,chunk_headers,chunk_idats,chunk_iend)

if __name__ == "__main__":
  print " [+] png chunk zlib thingy"
  if len(sys.argv) != 2:
    print " [-] usage: %s [png]" % sys.argv[0]
    sys.exit(0)
  (hdr,headers,idats,end) = png_split_chunks_simple(sys.argv[1])
  bruteforce_partial(hdr+headers, [],idats, end)