/*
 *  upkg: tool for manipulating Unreal Tournament packages.
 *  Copyright (C) 2009 Nick Bowler
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation; either version 2 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program; if not, write to the Free Software
 *  Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "upkg.h"
#include "pack.h"

#define MIN(a, b) ((a) < (b) ? (a) : (b))

struct upkg_name {
	unsigned long flags;
	char *name;
};

struct upkg_export {
	const char *name;

	long package, class, super;
	unsigned long flags;
	unsigned long size, offset;
};

struct upkg_import {
	const char *class_package, *class_name, *object_name;
	long package;
};

struct upkg_private {
	FILE *f;

	struct upkg_name   *names;
	struct upkg_export *exports;
	struct upkg_import *imports;

	unsigned long name_offset, export_offset, import_offset;
	unsigned char guid[16];
};

/*
 * Decode the compact index format from the upkg.  This format is fucked.
 * Stores the result in *val and returns the number of input bytes read (or 0
 * if the input is invalid, in which case *val is undefined).
 */
static size_t decode_index(long *val, unsigned char *bytes, size_t n)
{
	size_t i = 0;

	*val = 0;
	while (i < MIN(n, 5)) {
		/*
		 * Least significant bytes are first, so we need to do this
		 * nonsense.
		 */
		long tmp = bytes[i] & (i == 0 ? 0x3f : 0x7f);

		if (i > 0) tmp <<= 6;
		if (i > 1) tmp <<= 7*(i-1);
		*val += tmp;

		if (!(bytes[i] & (i == 0 ? 0x40 : 0x80))) {
			i++;
			break;
		}

		i++;
	}

	if (i > MIN(n, 5) || n == 0)
		return 0;
	if (bytes[0] & 0x80)
		*val = -*val;
	return i;
}

static struct upkg *init_upkg(unsigned char hdr[static UPKG_HDR_SIZE])
{
	struct upkg *pkg;

	pkg = malloc(sizeof *pkg);
	if (!pkg) {
		return NULL;
	}

	pkg->priv = malloc(sizeof *pkg->priv);
	if (!pkg->priv) {
		free(pkg);
		return NULL;
	}

	pkg->version      = unpack_16_le(hdr+4);
	pkg->license      = unpack_16_le(hdr+6);
	pkg->flags        = unpack_32_le(hdr+8);
	pkg->name_count   = unpack_32_le(hdr+12);
	pkg->export_count = unpack_32_le(hdr+20);
	pkg->import_count = unpack_32_le(hdr+28);

	pkg->priv->name_offset   = unpack_32_le(hdr+16);
	pkg->priv->export_offset = unpack_32_le(hdr+24);
	pkg->priv->import_offset = unpack_32_le(hdr+32);

	return pkg;
}

static int pkg_init_names(struct upkg *pkg)
{
	size_t rc, len, nbuf = 0;
	unsigned long index = 0;
	char buf[512];

	if (fseek(pkg->priv->f, pkg->priv->name_offset, SEEK_SET) != 0)
		return -1;

	pkg->priv->names = malloc(pkg->name_count * sizeof *pkg->priv->names);
	if (!pkg->priv->names)
		return -1;

	while (index < pkg->name_count) {
		struct upkg_name *name = &pkg->priv->names[index];

		/* Read some data into buffer. */
		if (!feof(pkg->priv->f)) {
			rc = fread(buf+nbuf, 1, sizeof buf-nbuf, pkg->priv->f);
			if (rc == 0)
				goto err;
			nbuf += rc;
		}

		if (pkg->version >= 64) {
			len = buf[0];
			if (nbuf <= len + 4 || buf[len])
				goto err;
			name->name = malloc(len);
			if (!name->name)
				goto err;
			memcpy(name->name, buf+1, len);
			name->flags = unpack_32_le(buf+len+1);
			len += 4;

			nbuf -= len + 1;
			memmove(buf, buf+len+1, nbuf);
			index++;
		} else {
			char *c = memchr(buf, 0, nbuf);
			if (!c || nbuf <= c - buf + 5)
				goto err;
			len = c - buf + 1;
			name->name = malloc(len);
			if (!name->name)
				goto err;
			memcpy(name->name, buf, len);
			name->flags = unpack_32_le(buf+len);
			len += 4;

			nbuf -= len;
			memmove(buf, buf+len, nbuf);
			index++;
		}
	}

	return 0;
err:
	for (unsigned i = 0; i < index; i++)
		free(pkg->priv->names[i].name);
	free(pkg->priv->names);
	return -1;
}

static int pkg_init_exports(struct upkg *pkg)
{
	size_t rc, len, nbuf = 0;
	unsigned long index = 0;
	char buf[512];

	if (fseek(pkg->priv->f, pkg->priv->export_offset, SEEK_SET) != 0)
		return -1;

	pkg->priv->exports = malloc(pkg->export_count * sizeof *pkg->priv->exports);
	if (!pkg->priv->exports)
		return -1;

	while (index < pkg->export_count) {
		struct upkg_export *export = &pkg->priv->exports[index];
		long tmp;

		/* Read some data into buffer. */
		if (!feof(pkg->priv->f)) {
			rc = fread(buf+nbuf, 1, sizeof buf-nbuf, pkg->priv->f);
			if (rc == 0)
				goto err;
			nbuf += rc;
		}

		len = 0;
		rc = decode_index(&export->class, buf+len, nbuf-len);
		if (rc == 0) goto err;
		len += rc;

		rc = decode_index(&export->super, buf+len, nbuf-len);
		if (rc == 0) goto err;
		len += rc;

		if (nbuf-len < 4) goto err;
		export->package = unpack_32_le(buf+len);
		len += 4;

		rc = decode_index(&tmp, buf+len, nbuf-len);
		if (rc == 0 || tmp < 0 || tmp >= pkg->name_count) goto err;
		export->name = pkg->priv->names[tmp].name;
		len += rc;

		if (nbuf-len < 4) goto err;
		export->flags = unpack_32_le(buf+len);
		len += 4;

		rc = decode_index(&export->size, buf+len, nbuf-len);
		if (rc == 0) goto err;
		len += rc;

		if (export->size) {
			rc = decode_index(&export->offset, buf+len, nbuf-len);
			if (rc == 0) goto err;
			len += rc;
		}

		nbuf -= len;
		memmove(buf, buf+len, nbuf);
		index++;
	}

	return 0;
err:
	free(pkg->priv->exports);
	return -1;
}

static int pkg_init_imports(struct upkg *pkg)
{
	size_t rc, len, nbuf = 0;
	unsigned long index = 0;
	char buf[512];

	if (fseek(pkg->priv->f, pkg->priv->import_offset, SEEK_SET) != 0)
		return -1;

	pkg->priv->imports = malloc(pkg->import_count * sizeof *pkg->priv->imports);
	if (!pkg->priv->imports)
		return -1;

	while (index < pkg->import_count) {
		struct upkg_import *import = &pkg->priv->imports[index];
		long tmp;

		/* Read some data into buffer. */
		if (!feof(pkg->priv->f)) {
			rc = fread(buf+nbuf, 1, sizeof buf-nbuf, pkg->priv->f);
			if (rc == 0)
				goto err;
			nbuf += rc;
		}

		len = 0;
		rc = decode_index(&tmp, buf+len, nbuf-len);
		if (rc == 0 || len < 0 || len >= pkg->name_count) goto err;
		import->class_package = pkg->priv->names[tmp].name;
		len += rc;

		rc = decode_index(&tmp, buf+len, nbuf-len);
		if (rc == 0 || len < 0 || len >= pkg->name_count) goto err;
		import->class_name = pkg->priv->names[tmp].name;
		len += rc;

		if (nbuf-len < 4) goto err;
		import->package = unpack_32_le(buf+len);
		len += 4;

		rc = decode_index(&tmp, buf+len, nbuf-len);
		if (rc == 0 || len < 0 || len >= pkg->name_count) goto err;
		import->object_name = pkg->priv->names[tmp].name;
		len += rc;

		nbuf -= len;
		memmove(buf, buf+len, nbuf);
		index++;
	}

	return 0;
err:
	free(pkg->priv->imports);
	return -1;
}

struct upkg *upkg_fopen(const char *path)
{
	unsigned char hdr_buf[UPKG_HDR_SIZE];
	struct upkg *pkg;
	FILE *f;

	if (!(f = fopen(path, "rb"))) {
		return NULL;
	}
	if (fread(hdr_buf, sizeof hdr_buf, 1, f) != 1) {
		goto err1;
	}
	if (unpack_32_le(hdr_buf) != UPKG_HDR_MAGIC) {
		goto err1;
	}

	/* Initialize package structure. */
	pkg = init_upkg(hdr_buf);
	if (!pkg) {
		goto err1;
	}
	pkg->priv->f = f;

	if (pkg_init_names(pkg) != 0) {
		goto err2;
	}

	if (pkg_init_exports(pkg) != 0) {
		goto err3;
	}

	if (pkg_init_imports(pkg) != 0) {
		goto err4;
	}

	return pkg;
err4:
	free(pkg->priv->exports);
err3:
	for (unsigned i = 0; i < pkg->name_count; i++)
		free(pkg->priv->names[i].name);
	free(pkg->priv->names);
err2:
	free(pkg->priv);
	free(pkg);
err1:
	fclose(f);
	return NULL;
}

int upkg_close(struct upkg *pkg)
{
	int rc = 0;

	if (pkg->priv->f) {
		rc = fclose(pkg->priv->f);

		for (unsigned i = 0; i < pkg->name_count; i++) {
			free(pkg->priv->names[i].name);
		}
	}

	free(pkg->priv->imports);
	free(pkg->priv->exports);
	free(pkg->priv->names);
	free(pkg->priv);
	free(pkg);

	return rc;
}

const char *upkg_get_name(struct upkg *pkg, unsigned long idx)
{
	if (idx >= pkg->name_count)
		return 0;
	return pkg->priv->names[idx].name;
}