//gptrefresh.c
/*
* Copyright (c) 2007, Timothy Perfitt
* All rights reserved.
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*     * Redistributions of source code must retain the above copyright
*       notice, this list of conditions and the following disclaimer.
*     * Redistributions in binary form must reproduce the above copyright
*       notice, this list of conditions and the following disclaimer in the
*       documentation and/or other materials provided with the distribution.
*     * Neither the name of the twocanoes nor the
*       names of its contributors may be used to endorse or promote products
*       derived from this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY TIMOTHY PERFITT ``AS IS'' AND ANY
* EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
* WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL <copyright holder> BE LIABLE FOR ANY
* DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
* (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
* LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
* ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
* (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
* SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

#import <Foundation/Foundation.h>
#include <unistd.h>
#define SIGNATURE 0xaa55

int verbose_flag;
static const struct partition_map {
	int	type;
	char	guid[37];
} partition_maps[] = {
	{ 0x07, "EBD0A0A2-B9E5-4433-87C0-68B6B72699C7"},
	{ 0xAF, "48465300-0000-11AA-AA11-00306543ECAC"},
	{ 0x00, "00000000-0000-0000-0000-000000000000"},
	{ 0xEE, "C12A7328-F81F-11D2-BA4B-00A0C93EC93B"},
	{ 0x8e, "E6D6D379-F507-44C2-A23C-238F2A3DF928"},
	{ 0xa8, "55465300-0000-11AA-AA11-00306543ECAC"}
};	
typedef struct  {
	int cyl;
	int head;
	int sector;
} chs;
typedef struct  {
	uint64_t signature;
	uint32_t revision;
	uint32_t header_size;	
	uint32_t header_checksum;
	uint32_t reserved1;
	uint64_t primary_lba;
	uint64_t backup_lba;
	uint64_t first_usable_lba;
	uint64_t last_usable_lba;
	uint8_t disk_guid[16];
	uint64_t partition_entry_lba;
	uint32_t number_of_partitions;
	uint32_t size_of_partition_entries;
	uint32_t partition_entry_checksum;
	uint8_t reserved2[420];
	
	
} gpt_header;
typedef struct _bios_partition {
	
	uint8_t flag;
	uint8_t start_head;
	uint8_t start_sector;
	uint16_t start_track;
	uint8_t type;
	uint8_t end_head;
	uint8_t end_sector;
	uint16_t end_track;
	uint32_t lba;
	uint32_t lba_size;
} bios_partition;

typedef struct _MBR{
	int test;
	uint8_t dosboot[446];
	bios_partition partition_array[4];
	uint16_t signature;
} MBR;

typedef struct _gpt_partition_entry{
	uint8_t partition_type_guid[16];
	uint8_t unique_partition_guid[16];
	uint64_t starting_lba;
	uint64_t ending_lba;	
	uint64_t attribute_bits;	
	uint8_t partition_name[72];	
	char partition_type_string[37];
	
} gpt_partition_entry;


void calculate_chs(chs *input_chs,int input_lba){
	int temp;
	
	int hpc=255;  //heads per cylinder
	int spt=63;	
	input_chs->cyl = input_lba / (hpc * spt);
	temp = input_lba % (hpc * spt);
	input_chs->head = temp / spt;
	input_chs->sector = temp % spt + 1;
	
	if (input_chs->cyl>1023) {
		input_chs->cyl=1023;
	//	input_chs->head=254;
	//	input_chs->sector=63;
	}

}


int ComparePartitionTables(bios_partition partition1[4],bios_partition partition2[4]){
	int i;
	BOOL isEqual=YES; 
	for (i=0;i<4;i++) {
		if (partition1[i].lba!=partition2[i].lba) {printf("lba is different: %i %i\n",partition1[i].lba,partition2[i].lba); isEqual=NO;}
		if (partition1[i].lba_size!=partition2[i].lba_size) {printf("lba_size is different: %i %i\n",partition1[i].lba_size,partition2[i].lba_size); isEqual=NO;}
		if (partition1[i].start_head!=partition2[i].start_head) {printf("start_head is different: %i %i\n",partition1[i].start_head,partition2[i].start_head); isEqual=NO;}
		if (partition1[i].start_sector!=partition2[i].start_sector) {printf("start_sector is different: %i %i\n",partition1[i].start_sector,partition2[i].start_sector); isEqual=NO;}
		if (partition1[i].start_track!=partition2[i].start_track) {printf("start_track is different: %i %i\n",partition1[i].start_track,partition2[i].start_track); isEqual=NO;}
		if (partition1[i].end_head!=partition2[i].end_head) {printf("end_head is different: %i %i\n",partition1[i].end_head,partition2[i].end_head); isEqual=NO;}
		if (partition1[i].end_sector!=partition2[i].end_sector) {printf("end_sector is different: %i %i\n",partition1[i].end_sector,partition2[i].end_sector); isEqual=NO;}
		if (partition1[i].end_track!=partition2[i].end_track) {printf("end_track is different: %i %i\n",partition1[i].end_track,partition2[i].end_track); isEqual=NO;}
		if (partition1[i].type!=partition2[i].type) {printf("type is different: %02x %02x\n",partition1[i].type,partition2[i].type); isEqual=NO;}
		if (isEqual==NO) return -1;
	}
	return 0;
	
}
void ConvertToGUID(uint8_t array_of_bytes[37],char *output_string){
	int i,cnt;
	uint16_t new_array[16];
	i=0;
	for (cnt=3;cnt>-1;cnt--) {
		new_array[cnt]=array_of_bytes[i];
		i++;
	}
	
	for (cnt=5;cnt>3;cnt--) {
		new_array[cnt]=array_of_bytes[i];
		i++;
	}
	
	for (cnt=7;cnt>5;cnt--) {
		new_array[cnt]=array_of_bytes[i];
		i++;
	}
	
	for (cnt=8;cnt<16;cnt++) {
		new_array[cnt]=array_of_bytes[i];
		i++;
	}
	sprintf(output_string,"%02X%02X%02X%02X-%02X%02X-%02X%02X-%02X%02X-%02X%02X%02X%02X%02X%02X",new_array[0],new_array[1],new_array[2],new_array[3],new_array[4],new_array[5],new_array[6],
			new_array[7],new_array[8],new_array[9],new_array[10],new_array[11],new_array[12],new_array[13],new_array[14],new_array[15]);
	
}
void read_bios_partition_table(FILE *file, MBR *mbr, int *active_partition) {
	int j;
	uint16_t cx;
	
	mbr->signature=SIGNATURE;
	fread(mbr->dosboot,1,446,file);
	for (j=0;j<4;j++) {
		fread(&mbr->partition_array[j].flag,1,1,file);
		if (mbr->partition_array[j].flag == 0x80) *active_partition=j+1;
		fread(&mbr->partition_array[j].start_head,1,1,file);
		fread(&cx,1,2,file);
		
		mbr->partition_array[j].start_sector=cx&63;
		mbr->partition_array[j].start_track=((cx & 192) << 2)|(cx & 65280)>>8;
		
		fread(&mbr->partition_array[j].type,1,1,file);
		fread(&mbr->partition_array[j].end_head,1,1,file);
		fread(&cx,1,2,file);
		mbr->partition_array[j].end_sector=cx&63;
		mbr->partition_array[j].end_track=((cx & 192) << 2)|(cx & 65280)>>8;
		fread(&mbr->partition_array[j].lba,1,4,file);
		fread(&mbr->partition_array[j].lba_size,1,4,file);
		
	}
	fread(&mbr->signature,1,2,file);

}

void read_gpt_header (FILE *file,gpt_header *gpt_header) {
	fread(&gpt_header->signature,1,8,file);
	fread(&gpt_header->revision,1,4,file);
	fread(&gpt_header->header_size,1,4,file);
	fread(&gpt_header->header_checksum,1,4,file);
	fread(&gpt_header->reserved1,1,4,file);
	fread(&gpt_header->primary_lba,1,8,file);
	fread(&gpt_header->backup_lba,1,8,file);
	fread(&gpt_header->first_usable_lba,1,8,file);
	fread(&gpt_header->last_usable_lba,1,8,file);
	fread(&gpt_header->disk_guid,1,16,file);
	fread(&gpt_header->partition_entry_lba,1,8,file);
	fread(&gpt_header->number_of_partitions,1,4,file);
	fread(&gpt_header->size_of_partition_entries,1,4,file);	
	fread(&gpt_header->partition_entry_checksum,1,4,file);
	fread(&gpt_header->reserved2,1,420,file);
	
}
void read_gpt(FILE *file, MBR *new_mbr, gpt_partition_entry *partition_entry, gpt_header *gpt_header, int partition_to_flag) {
	char partition_name[73];
	long long lda_start;
	chs chs;
	
	
	int i,index=0;
	
	new_mbr->signature=SIGNATURE; 
	
	for (i=0;i<gpt_header->number_of_partitions;i++){
		fread(&partition_entry[i].partition_type_guid,1,16,file);	
		ConvertToGUID(partition_entry[i].partition_type_guid,partition_entry[i].partition_type_string);
		if (i<4) {
			new_mbr->partition_array[i].type='\x00';
			new_mbr->partition_array[i].lba=0;
			new_mbr->partition_array[i].lba_size=0;
			new_mbr->partition_array[i].start_head=0;
			new_mbr->partition_array[i].start_sector=0;
			new_mbr->partition_array[i].start_track=0;
			new_mbr->partition_array[i].end_head=0;
			new_mbr->partition_array[i].end_sector=0;
			new_mbr->partition_array[i].end_track=0;
			new_mbr->partition_array[i].lba_size=0;
			new_mbr->partition_array[i].flag=0;
		}
		if (strcmp(partition_entry[i].partition_type_string, "00000000-0000-0000-0000-000000000000")==0) {
			memcpy(partition_entry[i].unique_partition_guid,"\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0",16);
			fseek(file, 112, SEEK_CUR); //skip the rest
			continue; //go around again
		}
		if (partition_to_flag==i+1) new_mbr->partition_array[i].flag=0x80;
			
		index++;		
		fread(&partition_entry[i].unique_partition_guid,1,16,file);	
		fread(&partition_entry[i].starting_lba,1,8,file);	
		if (i<4) new_mbr->partition_array[i].lba=partition_entry[i].starting_lba;
		fread(&partition_entry[i].ending_lba,1,8,file);	
		if (i<4) new_mbr->partition_array[i].lba_size=partition_entry[i].ending_lba-partition_entry[i].starting_lba+1;
		fread(&partition_entry[i].attribute_bits,1,8,file);	
		fread(&partition_entry[i].partition_name,1,72,file);
		
		// strip unicode from UTF-8
		int k,j;
		for (k=0,j=0;k<72;k=k+2,j++) {
			partition_name[j]=partition_entry[i].partition_name[k];
		}
		
		if (strcmp(partition_entry[i].partition_type_string,"C12A7328-F81F-11D2-BA4B-00A0C93EC93B")==0) {  //EFI
			lda_start=1;
			if (i<4) {
				new_mbr->partition_array[i].lba=1;
				new_mbr->partition_array[i].lba_size=partition_entry[i].ending_lba-new_mbr->partition_array[i].lba+1;
			}
		}
		else lda_start=partition_entry[i].starting_lba;
		if (i<4) {
			for (k=0;k<6;k++) {
				if (strcmp(partition_maps[k].guid,partition_entry[i].partition_type_string)==0) {
					new_mbr->partition_array[i].type=partition_maps[k].type;
					break;
				}
				
			}
		}
		
		if (i<4) {
			calculate_chs(&chs,lda_start);
			new_mbr->partition_array[i].start_head=chs.head;
			new_mbr->partition_array[i].start_sector=chs.sector;
			new_mbr->partition_array[i].start_track=chs.cyl;
			
			calculate_chs(&chs,partition_entry[i].ending_lba);
			new_mbr->partition_array[i].end_head=chs.head;
			new_mbr->partition_array[i].end_sector=chs.sector;
			new_mbr->partition_array[i].end_track=chs.cyl;
		}
		
	}
	
}

void print_gpt(gpt_partition_entry *partition_entry){
	
	printf(" -------------------------------------- GPT ----------------------------------\n");
	printf("      start       size  index  contents\n");
	int index=1;
	int i;
	for (i=0;i<128;i++) {
		if (strcmp(partition_entry[i].partition_type_string, "00000000-0000-0000-0000-000000000000")!=0)
			printf("%10lli %11lli %6i  GPT part - %s\n",partition_entry[i].starting_lba,partition_entry[i].ending_lba,index++,partition_entry[i].partition_type_string);
		
	}
	//

}
void print_bios_partition_table(MBR *mbr){
	int i;
	printf("\n\n");
	printf(" ----------------------------- BIOS Partition Table--------------------------\n");
	printf(" #: id  cyl  hd sec -  cyl  hd sec [     start -       size]\n");
	printf("------------------------------------------------------------------------\n");
	for (i=0;i<4;i++) {
		if ( ((mbr->partition_array[i].flag&0x80)>>7)==1) printf("*");
		else printf(" ");
		printf("%1i: %02X %4li %3i %3i - %4li %3i %3i [%10i - %10i]\n",i+1,mbr->partition_array[i].type,mbr->partition_array[i].start_track,
			   mbr->partition_array[i].start_head,mbr->partition_array[i].start_sector,mbr->partition_array[i].end_track,mbr->partition_array[i].end_head,mbr->partition_array[i].end_sector,
			   mbr->partition_array[i].lba,mbr->partition_array[i].lba_size);
	}
	
}
void write_new_gpt(int force_flag, char *device, char *mbr_path,MBR *new_mbr){

	if (force_flag==0) {
		char input[2];
		printf("Are you sure you want to write out a new partition table? <y/n> ");
		scanf("%1s",input);
		if (input[0]!='y') {
			printf("Exiting without writing\n");
			exit(-1);
		}
		
	}
	int fd;
	fd=open(device, O_RDWR|O_EXLOCK);
	if (fd==-1) {
		if (force_flag==0) {
			char input[2];
			printf("Cannot get exclusive access to %s.  You'll need to reboot if you choose to continue.  Continue? <y/n> ",device);
			scanf("%1s",input);
			if (input[0]!='y') {
				printf("exiting without writing\n");
				exit(-1);
			}
		}
		fd=open(device, O_RDWR|O_SHLOCK);
		if (fd==-1) {
			perror("error while opening file to write");
			exit (-1);
		}
		
	}	
	int i;
	uint8_t block[512];
	for (i=0;i<512;i++) {
		block[i]=0;
	}
	FILE *mbr_file;
	if (mbr_path) {
		mbr_file=fopen(mbr_path, "r");
		if (mbr_file==NULL){
			perror("error opening mbr file.");
			exit(-1);
		}
		fread(&block,1,446,mbr_file);	
	}
	int block_position=446;
	int j;
	uint16_t cx;
	for (j=0;j<4;j++) {
		memcpy(&block[block_position++],&new_mbr->partition_array[j].flag,1);
		memcpy(&block[block_position],&new_mbr->partition_array[j].start_head,1);
		block_position++;
		cx=((new_mbr->partition_array[j].start_track<<8)|((new_mbr->partition_array[j].start_track & 768 ) >> 2) ) | new_mbr->partition_array[j].start_sector;
		memcpy(&block[block_position],&cx,2);
		block_position+=2;
		memcpy(&block[block_position],&new_mbr->partition_array[j].type,1);
		block_position++;
		memcpy(&block[block_position],&new_mbr->partition_array[j].end_head,1);
		block_position++;
		cx=((new_mbr->partition_array[j].end_track<<8)|((new_mbr->partition_array[j].end_track & 768 ) >> 2) ) | new_mbr->partition_array[j].end_sector;
		memcpy(&block[block_position],&cx,2);
		block_position+=2;
		memcpy(&block[block_position],&new_mbr->partition_array[j].lba,4);
		block_position+=4;
		memcpy(&block[block_position],&new_mbr->partition_array[j].lba_size,4);
		block_position+=4;
	}
	memcpy(&block[block_position],&new_mbr->signature,2);
	int breakcnt=1;
	for (i=446;i<512;i++) {
		if (breakcnt==2) {
			breakcnt=1;
		}
		else breakcnt++;
	}
	if (!mbr_path) {
		lseek(fd,446,SEEK_SET); //jump past MBR
		block_position=446;
	}
	else block_position=0;
	
	write (fd,&block[block_position],512-block_position);
		
	
}
void logger(char *line) {
	if (verbose_flag) {
		printf("VERBOSE: %s\n",line);
	}
}
void usage() {
	printf("gptrefresh\n");
	printf("Utility to write out BIOS partition table from information gathered from GPT.  Also can ");
	printf("write out Master Boot Record and active partition\n\n");
	
	printf("usage: gptrefresh [-vwf] [-m mbr_path] [-a flag_partition] disk\n");
	printf("	-v: verbose output\n");
	printf("	-w: write out BIOS partition table (does not write out Master Boot Record.  Use -m to specify Master Boot Record)\n");
	printf("	-f: force writing without any warning messages\n");
	printf("	-m: write out master boot record at mbr_path with partition table.  Does nothing if -w is not specified\n");
	printf("	-a: flag flag_partition as active partition (flag_partition must be 1, 2, 3, or 4)\n");
	printf("'disk' is of the form /dev/disk0.\n");


}

int main (int argc, char ** argv) {	
	
	int active_partition=0;
	int partition_to_flag=0;

	verbose_flag=0;
	int write_flag=0,force_flag=0;
	char device[30];
	int ch;
	int newflag;
	char *mbr_path=NULL;
	while ((ch = getopt(argc, argv, "a:m:hvwf")) != -1) {
		switch (ch) {
			case 'v':
				verbose_flag = 1;
				break;
			case 'w':
				write_flag = 1;
				break;
			case 'f':
				force_flag = 1;
				break;
			case 'a':
				newflag=atoi(optarg);
				if (newflag>0 && newflag<5) partition_to_flag=newflag;
				break;
			case 'm':
				mbr_path=optarg;
					break;
			case '?':
			case 'h':
			default:
				usage();
				exit(-1);
		}
	}
	argc -= optind;
	argv += optind;

	if (argc==0) {
		usage();
		exit(-1);
	}

	if (geteuid()!=0) {
		printf("%s must be run as root\n",argv[0]);
		exit (-1);
	}
	
	strncpy(device,argv[0],30);
	
	logger("Opening device");
	
	FILE *file=fopen(device, "r");
	if (file==NULL) {
		perror("Error while opening file");
		exit(-1);
	}
	
	logger("Reading BIOS partition table");
	MBR mbr;
	read_bios_partition_table(file,&mbr,&active_partition);
	
	logger("Reading GPT Header");
	gpt_header gpt_header;
	
	read_gpt_header(file,&gpt_header);

	
	if (gpt_header.signature != 0x5452415020494645LL) {
		printf("Invalid GPT signature. Make sure this disk uses the GPT partitioning scheme.  Signature: %08llx\n",gpt_header.signature);
		exit (-1);
	}

	fseek(file,gpt_header.partition_entry_lba*512,SEEK_SET);  //skip to LBA 1; we should already be there, but just making sure
	

	MBR new_mbr;  // the mbr created from the GPT
	gpt_partition_entry partition_entry[gpt_header.number_of_partitions];
	
	logger("Reading GPT");	
	if (partition_to_flag!=0) active_partition=partition_to_flag;
	read_gpt(file, &new_mbr, partition_entry,&gpt_header,active_partition);
	
	fclose(file);
	
	logger("printing GPT");
	print_gpt(partition_entry);

	logger("printing bios header");
	
	print_bios_partition_table(&mbr);	
	
	logger("comparing bios partition with what it should be");
	int res=ComparePartitionTables(new_mbr.partition_array,mbr.partition_array);
	
	if (res==0) printf("The GPT and BIOS partitions appear to be consistant with each other.\n");
	else printf("The GPT and BIOS partitions appear to be inconsistant with each other!\n");
	
	if (write_flag==1) {
		logger("writing bios");
		printf("\n\n---------------------------------------------------------------\n");
		printf("-----------------New Partition Table to write------------------\n");
		printf("---------------------------------------------------------------\n");
		print_bios_partition_table(&new_mbr);	
		write_new_gpt(force_flag, device, mbr_path,&new_mbr);
	}
	printf ("\n");
    return 0;
}