png_uncorrupt.py

#!/usr/bin/env python

import re
import sys
import itertools
import collections

from struct import pack, pack_into, unpack, unpack_from, calcsize
from binascii import crc32, hexlify, unhexlify
from pprint import pprint

f = open(sys.argv[1])
fdat = f.read()

if len(sys.argv) == 4:
    # for running distributed
    offset = int(sys.argv[2])
    everyn = int(sys.argv[3])
else:
    offset = 0
    everyn = 1


out = open('out.png', 'w')

expected_types = ['sBIT', 'IHDR', 'pHYs', 'tEXt', 'IDAT', 'IEND']

chunk_idx = 0
chunks = []
# the chunk size has to have a leading zero
# because the file is less than 16MB in size
for m in re.finditer(r'\0.{0,3}[A-Za-z][A-Z][A-Z][A-Za-z]', fdat, re.M):
    end = m.end(0)
    start = end - 8
    ctype = fdat[start+4:start+8]
    clen_bytes = fdat[start:start+4]
    pcrc_bytes = fdat[start-4:start]
    # this doesn't happen on the challenge file, so don't worry about it
    if re.match(r'\x0a', clen_bytes, re.M):
        print 'WARNING: len might be bad'
    if re.match(r'\x0a', pcrc_bytes, re.M):
        print 'WARNING: pcrc might be bad'
    if ctype in expected_types:
        clen = unpack('!L', clen_bytes)[0]
        pcrc = unpack('!l', pcrc_bytes)[0]
        ccrc = unpack('!l', fdat[start+clen+8:start+clen+12])[0]
        print '%s %8u %8u %d %d' % (ctype, start+8, clen, pcrc, ccrc)
        chunks.append([ctype, start+8, clen, None, ccrc, chunk_idx])
        if chunk_idx > 0:
            plen = chunks[chunk_idx-1][2]
            poff = start - (chunks[chunk_idx-1][1] + 4)
            repl = plen - poff
            chunks[chunk_idx-1][3] = repl
            chunks[chunk_idx-1][4] = pcrc
        chunk_idx += 1

pprint(chunks)

def nth(iterable, n):
    "Returns every nth item"
    #return next(itertools.islice(iterable, 0, None, n), None)
    return itertools.islice(iterable, 0, None, n)

def consume(iterator, n):
    "Advance the iterator n-steps ahead. If n is none, consume entirely."
    # Use functions that consume iterators at C speed.
    if n is None:
        # feed the entire iterator into a zero-length deque
        collections.deque(iterator, maxlen=0)
    else:
        # advance to the empty slice starting at position n
        next(itertools.islice(iterator, n, n), None)

def addcr_at(data, pos):
    return data[0:pos] + '\x0d' + data[pos:]

def fix_data(data, pos_tup):
    for pos in sorted(pos_tup)[::-1]:
        data = addcr_at(data, pos)
    return data

# values from previous runs
known_errors = [
# 0-4
None, None, None, None, (553,),
# 5-9
(502, 26274, 125131), (71426,), None, (9618, 51473, 67850), (100443,),
# 10-14
(88347, 120996), None, (112266,), None, None,
# 15-19
None, None, None, None, None,
]

out.write('\x89PNG\x0d\x0a\x1a\x0a')

for idx in xrange(len(chunks)):
    chunk = chunks[idx]
    if chunk[2] > 0:
        data = fdat[chunk[1]-4:chunks[idx+1][1]-12]
    else:
        chunk[3] = 0
        data = chunk[0]
    if chunk[3] == 0:
        if crc32(data) == chunk[4]:
            print 'VALID %s %2u %11d' % (chunk[0], idx, chunk[4])
            out.write(pack('!L', (chunk[2])))
            out.write(data)
            out.write(pack('!l', (chunk[4])))
        else:
            print 'chunk correct length, but crc does not match :-('
            sys.exit(1)
    else:
        crc_target = chunk[4]
        maybe_errors = [m.start() for m in re.finditer(r'\x0a', data, re.M)]
        maybe_combos = itertools.combinations(maybe_errors, chunk[3])
        if False and known_errors[idx] is not None:
            maybe_combos = [known_errors[idx]]
        else:
            consume(maybe_combos, offset)
            maybe_combos = nth(maybe_combos, everyn)
            pass

        for combo in maybe_combos:
            maybe_fixed = fix_data(data, combo)
            new_crc = crc32(maybe_fixed)
            if crc_target == new_crc:
                print 'FIXED %s %2u %11d errors were %s' % (chunk[0], idx, chunk[4], repr(combo))
                out.write(pack('!L', (chunk[2])))
                out.write(maybe_fixed)
                out.write(pack('!l', (chunk[4])))
                break