#include <sys/mman.h>
#include <sys/stat.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <errno.h>
#include <fcntl.h>
#include <string.h>

#define END_ELF "\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF"

static unsigned char *mapfile(const char *path, unsigned long *size)
{
  int fd;
  struct stat s;
  unsigned char *ret;

  if ((fd = open(path, O_RDONLY, 0400)) < 0)
    {
      fprintf(stderr, "[X] %s : %s\n", path, strerror(errno));
      return NULL;
    }

  if (fstat(fd, &s) < 0)
    {
      perror("[X] fstat");
      return NULL;
    }
  *size = (unsigned long) s.st_size;
  if ((ret = (unsigned char *) mmap(0, s.st_size, PROT_READ, MAP_SHARED, fd, 0)) == NULL)
    {
      perror("[X] mmap");
      return NULL;
    }
  printf("[*] File %s mapped at %p, %ld bytes\n", path, ret, *size);
  return ret;
}

static int extract(const char *out, unsigned char *map, unsigned long size)
{
  unsigned char *cur_map;
  unsigned long end_map = (unsigned long) map + size;
  unsigned long cur_size;
  unsigned long n_ff, n_extra;
  int fd;

  if ((fd = open(out, O_RDWR | O_CREAT | O_TRUNC, 0644)) < 0)
    {
      perror(out);
      return 1;
    }
  
  for (cur_map = map, cur_size = n_ff = n_extra = 0; 
       (unsigned long) cur_map < end_map; )
    {
      printf("[*] Remaining : %ld bytes\r", 
             end_map - (unsigned long) cur_map);
      /* Skip FF's */
      if (!memcmp(END_ELF, cur_map, 16))
        {
          n_ff++;
          cur_size = 0;
          cur_map += 16;
          continue;
        }

      /* Skip extra data */
      if (cur_size == 2048)
        {
          n_extra++;
          cur_map += 64;
          cur_size = 0;
        }

      if (write(fd, cur_map, 16) != 16)
        {
          perror("write");
          return 1;
        }
      
      cur_size += 16;
      cur_map += 16;
    }
  printf("\n[*] All OK.\n[*] Stats: skipped %ld bytes of FF's\n" \
         "[*]        skipped %ld bytes of NAND extra buffer\n",
         n_ff * 16, n_extra * 64);
  return 0;
}

int main(int argc, const char **argv)
{
  unsigned char *mapped;
  unsigned long size;

  if (argc < 3)
    {
      fprintf(stderr, "Usage: %s <in-file> <out-file>\n", argv[0]);
      return 1;
    }
  if ((mapped = mapfile(argv[1], &size)) == NULL)
    return 1;
  return extract(argv[2], mapped, size);
}