/* Copyright 2013 Jason A. Donenfeld <Jason@zx2c4.com>. All Rights Reserved. */

#define _BSD_SOURCE
#define _GNU_SOURCE
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <stdint.h>
#include <stdbool.h>
#include <stddef.h>
#include <endian.h>
#include <fcntl.h>
#include <unistd.h>
#include <zlib.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/mman.h>

#define RG_MAGIC 0xFEEDBABE
#define HEADER_SIZE ( sizeof(uint32_t) * 5 + 0x80 )
struct File {
	uint32_t magic;			/* FEEDBABE */
	uint32_t size;			/* Length of file excluding header */
	uint32_t checksum;		/* 32-bit sum of all bytes in file and header, excluding checksum */
	uint32_t counter;		/* Unknown */
	uint32_t start_offset;		/* Unknown */
	char name[0x80];		/* Filename */
	uint8_t *data;			/* Pointer to mmap'd or allocated data of file */
	uint32_t decompressed_size;	/* Decompressed length of file */
	uint8_t *decompressed_data;	/* Pointer to allocated decompressed data */
};


uint8_t* decompress_data(uint8_t *data, uint32_t compressed_length, uint32_t decompressed_length)
{
	uint8_t *decompressed_data;
	z_stream strm;

	decompressed_data = malloc(decompressed_length);
	if (!decompressed_data)
		return NULL;

	strm.zalloc = Z_NULL;
	strm.zfree = Z_NULL;
	strm.opaque = Z_NULL;
	strm.avail_in = compressed_length;
	strm.next_in = data;
	strm.avail_out = decompressed_length;
	strm.next_out = decompressed_data;

	if (inflateInit(&strm) != Z_OK) {
		fprintf(stderr, "inflateInit: %s\n", strm.msg);
		goto error;
	}
	if (inflate(&strm, Z_SYNC_FLUSH) != Z_STREAM_END) {
		fprintf(stderr, "inflate: %s\n", strm.msg);
		goto error;
	}

	goto out;
error:
	free(decompressed_data);
	decompressed_data = NULL;
out:
	inflateEnd(&strm);
	return decompressed_data;
}

void free_file(struct File *file)
{
	if (file->decompressed_data)
		free(file->decompressed_data);
}

uint32_t checksum(uint8_t *data, uint32_t length)
{
	uint32_t sum = 0;

	while (length--) {
		/* The checksum does not include itself, so we skip that field. */
		if (length >= 8 && length < 12)
			continue;
		sum += data[length];
	}

	return sum;
}

bool read_file(uint8_t *data, struct File *file)
{
	uint32_t *packed_integers = (uint32_t *)data;
	file->magic = be32toh(packed_integers[0]);
	file->size = be32toh(packed_integers[1]);
	file->checksum = be32toh(packed_integers[2]);
	file->counter = be32toh(packed_integers[3]);
	file->start_offset = be32toh(packed_integers[4]);
	memcpy(&file->name, &packed_integers[5], sizeof(file->name));
	file->name[0x7f] = '\0';
	file->data = data + HEADER_SIZE;
	file->decompressed_data = NULL;
	file->decompressed_size = 0;
	if (!strcmp(file->name, "rg_conf")) {	
		/* When the file is gzipped, the length of the decompressed file is
		 * the first four bytes in big endian. */
		file->decompressed_size = be32toh(*((uint32_t *)file->data));
		file->decompressed_data = decompress_data(file->data + sizeof(uint32_t), file->size - sizeof(uint32_t), file->decompressed_size);
		if (!file->decompressed_data)
			file->decompressed_size = 0;
	}

	return file->magic == RG_MAGIC && file->checksum == checksum(data, HEADER_SIZE + file->size);
}

void print_file(struct File *file, bool valid)
{
	printf("==== %s (%s) ====\n", file->name, valid ? "Valid Checksum" : "INVALID CHECKSUM");
	printf("Length: %u bytes\n", file->size);
	if (file->decompressed_size)
		printf("Decompressed Length: %u bytes\n", file->decompressed_size);
}

void save_file(struct File *file, bool valid, uint32_t offset)
{
	char filename[0x100];
	char *slash;
	int fd;
	
	sprintf(filename, "%#.8x-%s%s", offset, file->name, valid ? "" : ".corrupt");
	for (slash = filename; (slash = strchr(slash, '/')) != NULL; ++slash)
		*slash = '_';

	printf("Saving to: %s\n", filename);

	fd = creat(filename, S_IRUSR | S_IWUSR);
	if (fd < 0) {
		perror("creat");
		exit(EXIT_FAILURE);
	}
	if (file->decompressed_size) {
		if (write(fd, file->decompressed_data, file->decompressed_size) < 0) {
			perror("write");
			exit(EXIT_FAILURE);
		}
	} else {
		if (write(fd, file->data, file->size) < 0) {
			perror("write");
			exit(EXIT_FAILURE);
		}
	}
	close(fd);
}

int main(int argc, char *argv[])
{
	int fd;
	bool valid;
	uint8_t *data, *next;
	struct stat sbuf;
	uint32_t needle;
	struct File file;

	if (argc != 2) {
		fprintf(stderr, "Usage: %s image_filename\n", argv[0]);
		return EXIT_FAILURE;
	}

	fd = open(argv[1], 0, O_RDONLY);
	if (fd < 0) {
		perror("open");
		return EXIT_FAILURE;
	}

	if (fstat(fd, &sbuf)) {
		perror("fstat");
		return EXIT_FAILURE;
	}

	data = mmap(NULL, sbuf.st_size, PROT_READ, MAP_SHARED, fd, 0);
	if (data == MAP_FAILED) {
		perror("mmap");
		return EXIT_FAILURE;
	}

	needle = htobe32(RG_MAGIC);
	next = data;
	for (;;) {
		next = memmem(next, sbuf.st_size - (next - data), &needle, sizeof(needle));
		if (!next)
			break;
		
		valid = read_file(next, &file);
		print_file(&file, valid);
		save_file(&file, valid, (uint32_t)(next - data));
		free_file(&file);

		printf("\n");

		++next;
	}

	munmap(data, sbuf.st_size);
	close(fd);
	return EXIT_SUCCESS;

}