idatfix - solution of fluxfingers for that png forensics challenge at pctf

idatfix.c

#define _GNU_SOURCE
#include <jh.h>
#include <string.h>
#include <arpa/inet.h>
#include <assert.h>
#include <fcntl.h>

   /* Table of CRCs of all 8-bit messages. */
   unsigned long crc_table[256];
   
   /* Flag: has the table been computed? Initially false. */
   int crc_table_computed = 0;
   
   /* Make the table for a fast CRC. */
   void make_crc_table(void)
   {
     unsigned long c;
     int n, k;
   
     for (n = 0; n < 256; n++) {
       c = (unsigned long) n;
       for (k = 0; k < 8; k++) {
         if (c & 1)
           c = 0xedb88320L ^ (c >> 1);
         else
           c = c >> 1;
       }
       crc_table[n] = c;
     }
     crc_table_computed = 1;
   }
   
   /* Update a running CRC with the bytes buf[0..len-1]--the CRC
      should be initialized to all 1's, and the transmitted value
      is the 1's complement of the final running CRC (see the
      crc() routine below)). */
   
   unsigned long update_crc(unsigned long crc, unsigned char *buf,
                            int len)
   {
     unsigned long c = crc;
     int n;

     for (n = 0; n < len; n++) {
       c = crc_table[(c ^ buf[n]) & 0xff] ^ (c >> 8);
     }
     return c;
   }
   
   /* Return the CRC of the bytes buf[0..len-1]. */
   unsigned long crc(unsigned char *buf, int len)
   {
     return update_crc(0xffffffffL, buf, len) ^ 0xffffffffL;
   }

bool fixidat(char *data, int data_len, int tofix, uint32_t expected_crc, int offset, bool toplevel) {
  char saved[data_len+tofix];
  memcpy(saved, data, data_len+tofix);
  if (tofix == 0) {
    return crc(data-4, data_len+4) == expected_crc;
  }
  if (toplevel) {
    int chars = count_char_occurences_in_buf(saved, data_len, 0x0a);
    for (int i=0; i<chars; i++) printf("X");
    puts("");
  }
  int j=0;
  for (int i=0; i<=data_len; i++) {
    if (saved[i] != 0x0a) continue;
    if (toplevel) printf(".");
    memcpy(data, saved, i);
    data[i] = 0xd;
    memcpy(data+i+1, saved+i, data_len+tofix-i-1);
    if (fixidat(data, data_len+1, tofix-1, expected_crc, i+2, false)) return true;
  }
  return false;
}

int main(void) {

  make_crc_table();
  setbuf(stdout, NULL);
  size_t len;
  char *data = slurp_file("fixed_partial_393365.png", &len, 0);
  data = realloc(data, len + 1024);
  char *dataend = data + len;
  for (int i=393366; i<len; i++) {
    if (strncmp(data+i, "IDAT", 4) == 0) {
      char *nextidat = memmem(data+i+1, len-i-1, "IDAT", 4);
      if (nextidat == NULL) nextidat = memmem(data+i+1, len-i-1, "IEND", 4);
      assert(nextidat);
      uint32_t header_len = ntohl(*(uint32_t*)(data+i-4));
      int real_len = nextidat - (data+i) - 12;
      printf("IDAT len=%u, real_len=%d\n", header_len, real_len);
      assert(real_len <= header_len);
      int tofix = header_len - real_len;
      printf("tofix=%d\n", tofix);
      uint32_t crcval = ntohl(*(uint32_t*)(nextidat-8));
      memmove(nextidat-8+tofix, nextidat-8, dataend-(nextidat-8));
      len += tofix;
      dataend = data + len;
      assert(fixidat(data+i+4, real_len, tofix, crcval, 0, true) == true);
      puts("");

      char name[200];
      sprintf(name, "fixed_partial_%d.png", i);
      write_file(name, data, len, O_EXCL|O_CREAT);
    }
  }
  return 0;
}