/*
 * ocalc.c - utility to calculate the instruction offset inside a mach-o binary
 * coded by ghalen <ghalen@hack.se> -- 2009/03/20
 *
 * This code is roughly based on 'offset12pl' by Fractal Guru <reverse@put.as>
 * Remember to use the right offset for the arch you're playing with :)
 * 
 * - ghalen / http://kmem.se
 */

#include <mach-o/loader.h>
#include <mach-o/nlist.h>
#include <mach-o/fat.h>

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

void
get_offset(const char *buf, uint32_t base, uint32_t my_offset, uint32_t arch)
{
	int i, j;
	uint32_t addr = 0x0, size = 0x0, offset = 0x0, real_offset = 0x0;
	uint32_t ncmds, nsects, cmdsize, cmd;

	struct mach_header *mh; /* mach header */
	struct load_command *lc; /* load commands */
	struct segment_command *sc; /* segment command */
	struct section *st; /* section */

	mh = (struct mach_header *)(buf + base);
	lc = (struct load_command*)((char*)mh + sizeof(struct mach_header));

	ncmds = (arch == CPU_TYPE_I386)?(mh->ncmds):(ntohl(mh->ncmds));
	
	for (i = 0; i < ncmds; i++) {
		cmd = (arch == CPU_TYPE_I386)?(lc->cmd):(ntohl(lc->cmd));
		if (cmd != LC_SEGMENT)
			goto end;

		sc = (struct segment_command *)lc;

		if (strcmp(sc->segname, "__TEXT"))
			goto end;

		st = (struct section *)((char *)sc + sizeof(struct segment_command));

		nsects = (arch == CPU_TYPE_I386)?(sc->nsects):(ntohl(sc->nsects));
		for (j = 0; j < nsects; j++) {
			if (!strcmp(st[j].sectname, "__text"))
				break;
		}

		addr = (arch == CPU_TYPE_I386)?(st[j].addr):(ntohl(st[j].addr));
		size = (arch == CPU_TYPE_I386)?(st[j].size):(ntohl(st[j].size));
		offset = (arch == CPU_TYPE_I386)?(st[j].offset):(ntohl(st[j].offset));

end:
		cmdsize = (arch == CPU_TYPE_I386)?(lc->cmdsize):(ntohl(lc->cmdsize));
		lc = (struct load_command*)((char*)lc + cmdsize);
	}

	if (my_offset < addr || my_offset > (addr + size)) {
		fprintf(stderr, " ! your offset is outside the code region for this arch\n");
		return;
	}

	real_offset = base + offset + my_offset - addr;
	fprintf(stderr, " -> actual offset for this arch: %d (%p)\n", real_offset, (void *)real_offset);
}

int
main(int argc, char *argv[])
{
	unsigned int sz;
	uint32_t intel_base = 0x0, my_offset, narchs = 0x0, ppc_base = 0x0;
	int i;
	char *buf;
	FILE *fp;

	struct fat_header *fh; /* fat header */
	struct fat_arch *fa; /* fat arch */
	struct mach_header *mh; /* mach header */

	if (argc != 3) {
		fprintf(stderr, "usage: %s <file> <offset>\n", argv[0]);
		return 0;
	}

	my_offset = (uint32_t)strtoul(argv[2], NULL, 0);

	fp = fopen(argv[1], "r");
	if (!fp) {
		fprintf(stderr, "could not open file\n");
		exit(EXIT_FAILURE);
	}

	fseek(fp, 0, SEEK_END);
	sz = ftell(fp);
	fseek(fp, 0, SEEK_SET);

	buf = malloc(sz);
	fread(buf, sz, 1, fp);
	fclose(fp);

	/* do we have a fat binary?
	 * if that's the case, start parsing the fat header
	 */
	if (ntohl(*(unsigned int*)buf) == FAT_MAGIC) {
		fh = (struct fat_header *)(buf);
		narchs = ntohl(fh->nfat_arch);
		fprintf(stderr, " -> found fat header with %x archs\n", narchs);

		fa = (struct fat_arch *)(buf + sizeof(struct fat_header));
		for (i = 0; i < narchs; i++) {
			switch (ntohl(fa->cputype)) {
			case CPU_TYPE_I386:
				intel_base = ntohl(fa->offset);
				fprintf(stderr, " + intel offset: %x\n", intel_base);
				get_offset(buf, intel_base, my_offset, CPU_TYPE_I386);
				break;
			case CPU_TYPE_POWERPC:
				ppc_base = ntohl(fa->offset);
				fprintf(stderr, " + ppc offset: %x\n", ppc_base);
				get_offset(buf, ppc_base, my_offset, CPU_TYPE_POWERPC);
				break;
			default:
				break;
			}

			fa += ((sizeof(struct fat_arch) * i) + 1);
		}
	} else { /* if not, check the mach header */
		fprintf(stderr, " -> this is not a fat binary, checking arch ... ");
		fflush(stderr);
		
		mh = (struct mach_header *)(buf);
		switch (mh->cputype) {
		case CPU_TYPE_I386:
			fprintf(stderr, "i386 found\n");
			get_offset(buf, 0, my_offset, CPU_TYPE_I386);
			break;
		case CPU_TYPE_POWERPC:
			fprintf(stderr, "ppc found\n");
			get_offset(buf, 0, my_offset, CPU_TYPE_POWERPC);
			break;
		default:
			fprintf(stderr, "unknown\nthis is binary is neither I386 nor PPC.\n");
			break;
		}
	}

	free(buf);
	return 0;
}