#import "SegmentManager.h"
#import <stdlib.h>
#import <string.h>
#import <mach-o/loader.h>
#import <sys/types.h>
#import <sys/stat.h>
#import <mach-o/ldsyms.h>
#import <libc.h>
#import <objc/maptable.h>
#import <mach-o/fat.h>

#define REALLY_SWAP_LONG(a) ( ((a) << 24) | \
		      (((a) << 8) & 0x00ff0000) | \
		      (((a) >> 8) & 0x0000ff00) | \
	((unsigned long)(a) >> 24) )

#define SWAP_LONG(a) (shouldSwap ? REALLY_SWAP_LONG(a) : (a))

@implementation SegmentManager

static NXMapTable *managers = NULL;

+initialize
{
    static BOOL beenHere = NO;
    
    if (!beenHere) {
	beenHere = YES;
	managers = NXCreateMapTable(NXStrValueMapPrototype, 0);
    }
    return self;
}

-init
{
    relocSize = sizeof(Segment);
    rmFlags.shouldSortRelocs = YES;
    [self readInAllRelocs];
    lastReloc = NULL;
    if (numRelocs) {
	if (images[0].name)
	    NXMapInsert(managers, images[0].name, self);
	return self;
    } else {
	[self free];
	return nil;
    }
}

-(void)setCpuType: (cpu_type_t)type
{
    static cpu_type_t defaultCpuType = 0;
    if (type)
	cpuType = type;
    else {
	if (!defaultCpuType) {
	    struct host_basic_info host_basic_info;
	    unsigned int count = HOST_BASIC_INFO_COUNT;
	    
	    if (host_info(host_self(), HOST_BASIC_INFO,
			  (host_info_t)(&host_basic_info),
			  &count) == KERN_SUCCESS)
		defaultCpuType = host_basic_info.cpu_type;
	}
	cpuType = defaultCpuType;
    }
}

-initFile: (STR)theName skipShlibs: (BOOL)shouldSkipShlibs withDesc: (int)desc cpuType: (cpu_type_t)type
{
    SegmentManager *oldManager;
    if (oldManager = NXMapGet(managers, theName)) {
	close(desc);
	[self free];
	return oldManager;
    } else  {
	name = malloc(strlen(theName) + 1);
	strcpy(name, theName);
	skipShlibs = shouldSkipShlibs;
	imageFd = desc;
	[self setCpuType: type];
	[self getImages];
	return [self init];
    }
}

static const char loadedName[] = "loaded file";

-(struct mach_header *)headerFromStart: (void *)start
{
    struct mach_header *header = start;
    if (header->magic == MH_MAGIC && header->cputype == cpuType)
	return header;
    else {
	struct fat_header *fatHeader = start;
	if ((fatHeader->magic == FAT_MAGIC)
	    || (fatHeader->magic == REALLY_SWAP_LONG(FAT_MAGIC))) {
	    BOOL shouldSwap = !(fatHeader->magic == FAT_MAGIC);
	    unsigned int nArches;
	    struct fat_arch *fatArch, *foundArch;
	    
	    for (nArches = SWAP_LONG(fatHeader->nfat_arch),
		 fatArch = (struct fat_arch *)(fatHeader + 1),
		 foundArch = NULL;
		 nArches && !foundArch;
		 fatArch++, nArches--) {
		if (cpuType == SWAP_LONG(fatArch->cputype))
		    foundArch = fatArch;
	    }
	    if (foundArch)
		return start + SWAP_LONG(foundArch->offset);
	    else
		return NULL;
	} else
	    return NULL;
    }
}

-initHeader: (struct mach_header *)header withSize: (int)size andSlide:(unsigned long)slide
{
    numImages = 1;
    [self setCpuType: 0];
    images = malloc(sizeof(*images));
    images->header = [self headerFromStart: header];
    images->size = size;
    images->name = NULL;
    images->mtime = 0;
    images->slide = slide;
    name = malloc(sizeof(loadedName));
    strcpy(name, loadedName);
    return [self init];
}

// Assume header and all of it's load commands have been 
// warped onto the current memory space.

-initHeader: (struct mach_header *)header withSlide:(unsigned long)slide
{
  int i;
  struct load_command* cmd;
  unsigned long maxAddr;
  unsigned long hAddr;

  // Compute the size of the module from the load commands.

  [self setCpuType: 0];
  maxAddr = hAddr = (unsigned) header = [self headerFromStart:header];

  for (cmd = (struct load_command *)(header + 1), i = 0;
       i < header->ncmds;
       ((void *)cmd) += cmd->cmdsize, i++) {

    unsigned long topAddr;

    if (cmd->cmd == LC_SEGMENT)
      {
	struct segment_command *segCmd = (struct segment_command *)cmd;
	topAddr = segCmd->vmaddr + segCmd->vmsize + slide;

	if (topAddr > maxAddr)
	  maxAddr = topAddr;
      }
  }

  return [self initHeader: header 
	         withSize: maxAddr - (unsigned)header
	         andSlide: slide];
}

+newHeader:(struct mach_header*)header withSlide:(unsigned long)slide
{
  return [[self alloc] initHeader:header withSlide:slide];
}

-initHeader: (struct mach_header *)header withSize: (int)size
{
  return [self initHeader:header withSize:size andSlide:0];
}

-initHeader: (struct mach_header *)header
{
  return [self initHeader:header withSlide:0];
}

-initFile: (STR)theName
{
    return [self initFile: theName skipShlibs: NO withDesc: -1 cpuType: 0];
}

+newExecutable: (STR)theName skipShlibs: (BOOL)shouldSkipShlibs withDesc: (int)desc cpuType: (cpu_type_t)type
{
    SegmentManager *sm = [[super alloc] initFile: (STR)theName
                                    skipShlibs: shouldSkipShlibs
				      withDesc: desc
				   cpuType: type];
    if (sm && [sm isExecutable])
	 return sm;
    else {
	[sm free];
	return nil;
    }
}

+newExecutable: (STR)theName
{
    return [self newExecutable: theName skipShlibs: NO withDesc: -1 cpuType: 0];
}

+newCore: (STR)theName
{
    SegmentManager *sm = [[super alloc] initFile: (STR)theName];
    if (sm && [sm isCore])
	 return sm;
    else {
	[sm free];
	return nil;
    }
}

+newFile: (STR)theName
{
    return [[super alloc] initFile: (STR)theName];
}

+newShlib: (STR)theName cpuType: (cpu_type_t)type
{
    SegmentManager *sm = [[super alloc] initFile: (STR)theName
                                    skipShlibs: YES
				      withDesc: -1
				       cpuType: type];
    if (sm && [sm isShlib])
	 return sm;
    else {
	[sm free];
	return nil;
    }
}

+newShlib: (STR)theName
{
    return [self newShlib: theName cpuType: 0];
}

+newHeader: (struct mach_header *)header withSize: (int)size
{
    return [[super alloc] initHeader: header withSize: size];
}

-(STR)executableName
{
    return name;
}

-(long)mtime
{
    return images[0].mtime;
}

-free
{
    NXMapRemove(managers, name);
    free(name); name = NULL;
    return [super free];
}

-invalidate
{
    int count;
    Image *image;
    
    if (images) {
	for (count = numImages, image = images; count; count--, image++) {
	    if (image->deallocate)
		vm_deallocate(task_self(),
				(vm_address_t)image->header,
				image->size);
	}
	free(images); images = NULL;
	numImages = 0;
    }
    return [super invalidate];
}

-(void)getImages
{
    struct stat imageStat;
    struct load_command *loadCmd;
    struct fvmlib_command *fvmCmd;
    struct mach_header *header;
    int i;
    Image *image;
    numImages = 0;
    if (imageFd < 0)
	imageFd = open(name, O_RDONLY, 0);
    if (imageFd >= 0) {
	extern int errno;
	errno = 0;
	fstat(imageFd, &imageStat);
	if (map_fd(imageFd,
		   0,
		   (vm_address_t *)&header,
		   YES,
		   imageStat.st_size) == 0) {
//	    printf("Map_fp successfull.\n");
	    if (header && (header = [self headerFromStart: header])) {
//	        printf("Got header.\n");
		numImages = 1;
		if (skipShlibs) {
		    images = malloc(sizeof(Image));
		    images->name = name;
		    images->header = header;
		    images->size = imageStat.st_size;
		    images->deallocate = YES;
		    images->slide = 0;
		    images->mtime = imageStat.st_mtime;
		} else {
		    for (i = 0, loadCmd = (struct load_command *)(header + 1);
			 i < header->ncmds;
			 i++, ((void *)loadCmd) += loadCmd->cmdsize) {
			if (loadCmd->cmd == LC_LOADFVMLIB) 
			    numImages++;
		    }
		    image
		    = images
		    = calloc(numImages, sizeof(Image));
		    image->name = name;
		    image->header = header;
		    image->size = imageStat.st_size;
		    image->deallocate = YES;
		    image->slide = 0;
		    image->mtime = imageStat.st_mtime;
		    image++;
		    for (i = 0, loadCmd = (struct load_command *)(header + 1);
			 i < header->ncmds;
			 i++, ((void *)loadCmd) += loadCmd->cmdsize) {
			int shlibFd;
			if (loadCmd->cmd == LC_LOADFVMLIB) {
			    fvmCmd = (struct fvmlib_command *)loadCmd;
			    image->name
			    = (void *)fvmCmd + fvmCmd->fvmlib.name.offset;
			    shlibFd = open(image->name, O_RDONLY, 0);
			    if (shlibFd >= 0) {
				fstat(shlibFd, &imageStat);
				image->size = imageStat.st_size;
				image->mtime = imageStat.st_mtime;
				if (map_fd(shlibFd,
					   0,
					   (vm_address_t *)&image->header,
					   YES,
					   image->size)
					  == 0) {
			            image->header
				    = [self headerFromStart: image->header];
				    image->deallocate = YES;
				    image++;
				} else
				    numImages--;
				close(shlibFd);
			    } else
				numImages--;
			}
		    }
		}
	    } else {
		printf("Error: file had bad header.\n");
	    }
	} else {
//	    printf("Map_fd failed, errno = %d.\n", errno);
	}
	close(imageFd);
	imageFd = -1;
    }
}
						    
-(int)numSegments
{
    int count, numSegs = 0, i;
    Image *image;
    struct load_command *loadCmd;
    struct mach_header *header;
    for (count = numImages,
	 image = images;
	 count;
	 count--, image++) {
	header = image->header;
	for (i = 0, loadCmd = (struct load_command *)(header + 1);
	     i < header->ncmds;
	     i++, ((void *)loadCmd) += loadCmd->cmdsize) {
	    if (loadCmd->cmd == LC_SEGMENT) 
		numSegs++;
	}
    }
    return numSegs;
}

-(BOOL)validate
{
    int nImages, i;
    Image *image;
    struct mach_header *header;
    struct load_command *loadCmd;
    struct segment_command *segCmd;
    BOOL allOK = YES;
    for (nImages = numImages,
	 image = images;
	 allOK && nImages;
	 nImages--, image++) {
	header = image->header;
	for (i = 0, loadCmd = (struct load_command *)(header + 1);
	     allOK && i < header->ncmds;
	     i++, ((void *)loadCmd) += loadCmd->cmdsize) {
	    if (loadCmd->cmd == LC_SEGMENT) {
		segCmd = (struct segment_command *)loadCmd;
		if ((segCmd->fileoff + segCmd->vmsize) > image->size)
		    allOK = NO;
	    }
	    if (!loadCmd->cmdsize)
		allOK = NO;
	}
    }
    return allOK;
}

-(void)readInSegments
{
    int nImages, i;
    Image *image;
    struct mach_header *header;
    struct load_command *loadCmd;
    struct segment_command *segCmd;
    Segment *segment = relocs;
    for (nImages = numImages,
	 image = images;
	 nImages;
	 nImages--, image++) {
	header = image->header;
	for (i = 0, loadCmd = (struct load_command *)(header + 1);
	     i < header->ncmds;
	     i++, ((void *)loadCmd) += loadCmd->cmdsize) {
	    if (loadCmd->cmd == LC_SEGMENT) {
		segCmd = (struct segment_command *)loadCmd;
		segment->address = segCmd->vmaddr + image->slide;
		segment->size = segCmd->vmsize;
		segment->data = (vm_address_t)header + segCmd->fileoff;
		segment->maxAddress = segment->address + segment->size;
		segment->displacement = segment->data - segment->address;
		segment->maxData = segment->data + segment->size;
		segment->segName = segCmd->segname;
		segment->rFlags.readIn = YES;
		((void *)segment) += relocSize;
	    }
	}
    }
}

-(void)_readInAllRelocs
{
    numRelocs = [self numSegments];
    relocs = calloc(numRelocs, relocSize);
    [self readInSegments];
    rmFlags.invalid = NO;
}

-(struct mach_header *)getMachHeader
{
    if (numImages)
	return images[0].header;
    else
	return NULL;
}

-(unsigned)getSlide
{
  if (numImages)
    return images[0].slide;
  else
    return 0;
}

-(int)getNumMachHeaders
{
    return numImages;
}

-(struct mach_header **)getMachHeaders
{
    struct mach_header **headers;
    int headerIndex;
    headers = malloc((numImages + 1) * sizeof(*headers));
    for (headerIndex = 0; headerIndex < numImages; headerIndex++)
	headers[headerIndex] = images[headerIndex].header;
    headers[numImages] = NULL;
    return headers;
}

-(unsigned*)getSlides
{
    unsigned *slides;
    int slideIndex;
    slides = malloc((numImages + 1) * sizeof(*slides));
    for (slideIndex = 0; slideIndex < numImages; slideIndex++)
	slides[slideIndex] = images[slideIndex].slide;
    slides[numImages] = 0;
    return slides;
}

-(struct mach_header **)getMachHeadersWithNames: (char ***)names
{
    struct mach_header **headers;
    char **theNames;
    int headerIndex;
    headers = malloc((numImages + 1) * sizeof(*headers));
    *names = theNames = malloc((numImages + 1) * sizeof(*theNames));
    for (headerIndex = 0; headerIndex < numImages; headerIndex++) {
	headers[headerIndex] = images[headerIndex].header;
	theNames[headerIndex] = images[headerIndex].name;
    }
    headers[numImages] = NULL;
    theNames[numImages] = NULL;
    return headers;
}

-(unsigned *)getSlidesWithNames: (char ***)names
{
    unsigned *slides;
    char **theNames;
    int slideIndex;
    slides = malloc((numImages + 1) * sizeof(*slides));
    *names = theNames = malloc((numImages + 1) * sizeof(*theNames));
    for (slideIndex = 0; slideIndex < numImages; slideIndex++) {
	slides[slideIndex] = images[slideIndex].slide;
	theNames[slideIndex] = images[slideIndex].name;
    }
    slides[numImages] = 0;
    theNames[numImages] = NULL;
    return slides;
}

-(struct load_command *)findLoadCommand: (unsigned long)command
                              forHeader: (struct mach_header *)header
{
    int i;
    struct load_command *loadCmd, *foundCmd = NULL;
    for (i = 0, loadCmd = (struct load_command *)(header + 1);
	 i < header->ncmds && !foundCmd;
	 i++, ((void *)loadCmd) += loadCmd->cmdsize) {
	if (loadCmd->cmd == command)
	    foundCmd = loadCmd;
    }
    return foundCmd;
}

-(int)numCommands: (unsigned long)command
        forHeader: (struct mach_header *)header
{
    int i, nCommands = 0;
    struct load_command *loadCmd;
    for (i = 0, loadCmd = (struct load_command *)(header + 1);
	 i < header->ncmds;
	 i++, ((void *)loadCmd) += loadCmd->cmdsize) {
	if (loadCmd->cmd == command)
	    nCommands++;
    }
    return nCommands;
}

-(struct symtab_command *)symCmd
{
    return (struct symtab_command *)
             [self findLoadCommand: LC_SYMTAB
	                 forHeader: images[0].header];
}

-(struct dysymtab_command *)dysymCmd
{
    return (struct dysymtab_command *)
             [self findLoadCommand: LC_DYSYMTAB
	                 forHeader: images[0].header];
}

-(char *)stringTable
{
    struct symtab_command *symCmd = [self symCmd];
    return (char *)(images[0].header) + symCmd->stroff;
}

-(struct nlist *)symbolTable
{
    struct symtab_command *symCmd = [self symCmd];
    return (struct nlist*)((char *)(images[0].header) + symCmd->symoff);
}

-(BOOL)isDylib
{
  return images && images[0].header->filetype == MH_DYLIB;
}

-(BOOL)isCore
{
    return images && images[0].header->filetype == MH_CORE;
}

-(BOOL)isShlib
{
    return images && images[0].header->filetype == MH_FVMLIB;
}

-(BOOL)isExecutable
{
    BOOL ret = images && ((images[0].header->filetype == MH_EXECUTE)
		      || (images[0].header->filetype == MH_OBJECT)
                      || (images[0].header->filetype == MH_PRELOAD));
//    if (ret)
//	printf("Is executable.\n");
//    else
//	printf("Isn't executable.\n");
    return ret;
}

-(BOOL)isSymbolStub:(void*)pc
{
  const struct section *sec;
  sec = [self getSeg:"__TEXT" sect:"__picsymbol_stub"];
  if (sec == 0)
    sec = [self getSeg:"__TEXT" sect:"__symbol_stub"];
  
  if (sec 
      && pc >= (void *) ((char *) 0 + (sec->addr + [self getSlide]))
      && pc <  (void *) ((char *) 0 + (sec->addr + [self getSlide] 
				                 + sec->size)))
    return 1;
  else
    return 0;
}

@end


