/***************************************************************************
 *             __________               __   ___.
 *   Open      \______   \ ____   ____ |  | _\_ |__   _______  ___
 *   Source     |       _//  _ \_/ ___\|  |/ /| __ \ /  _ \  \/  /
 *   Jukebox    |    |   (  <_> )  \___|    < | \_\ (  <_> > <  <
 *   Firmware   |____|_  /\____/ \___  >__|_ \|___  /\____/__/\_ \
 *                     \/            \/     \/    \/            \/
 * $Id: ipodpatcher.c 12237 2007-02-08 21:31:38Z dave $
 *
 * Copyright (C) 2006-2007 Dave Chapman
 *
 * All files in this archive are subject to the GNU General Public License.
 * See the file COPYING in the source tree root for full license agreement.
 *
 * This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY
 * KIND, either express or implied.
 *
 ****************************************************************************/

#include <stdio.h>
#include <unistd.h>
#include <fcntl.h>
#include <string.h>
#include <stdlib.h>
#include <inttypes.h>
#include <sys/types.h>
#include <sys/stat.h>

#include "ipodpatcher.h"
#include "ipodio.h"

#define VERSION "2.0 with v2.0 bootloaders"

int verbose = 0;

enum {
	NONE,
	LIST_IMAGES,
#ifdef WITH_BOOTOBJS
	INTERACTIVE,
	INSTALL,
#endif
	DELETE_BOOTLOADER,
	ADD_BOOTLOADER,
	READ_FIRMWARE,
	WRITE_FIRMWARE,
	READ_AUPD,
	WRITE_AUPD,
	READ_PARTITION,
	WRITE_PARTITION,
	FORMAT_PARTITION,
	CONVERT_TO_FAT32
};

void print_macpod_warning(void)
{
	printf("************************************************************************\n");
	printf("*** WARNING FOR ROCKBOX USERS\n");
	printf("*** You must convert this iPod to FAT32 format (aka a \"Winpod\")\n");
	printf("*** if you want to run Rockbox.  Rockbox WILL NOT work on this iPod.\n");
	printf("*** See http://www.rockbox.org/twiki/bin/view/Main/IpodConversionToFAT32\n");
	printf("************************************************************************\n");
}

void print_usage(void)
{
	fprintf(stderr,"Usage: ipodpatcher --scan\n");
#ifdef __WIN32__
	fprintf(stderr,"    or ipodpatcher [DISKNO] [action]\n");
#else
	fprintf(stderr,"    or ipodpatcher [device] [action]\n");
#endif
	fprintf(stderr,"\n");
	fprintf(stderr,"Where [action] is one of the following options:\n");
#ifdef WITH_BOOTOBJS
	fprintf(stderr,"        --install\n");
#endif
	fprintf(stderr,"  -l,   --list\n");
	fprintf(stderr,"  -r,   --read-partition     bootpartition.bin\n");
	fprintf(stderr,"  -w,   --write-partition    bootpartition.bin\n");
	fprintf(stderr,"  -rf,  --read-firmware      filename.ipod\n");
	fprintf(stderr,"  -rfb, --read-firmware-bin  filename.bin\n");
	fprintf(stderr,"  -wf,  --write-firmware     filename.ipod\n");
	fprintf(stderr,"  -wfb, --write-firmware-bin filename.bin\n");
#ifdef WITH_BOOTOBJS
	fprintf(stderr,"  -we,  --write-embedded\n");
#endif
	fprintf(stderr,"  -a,   --add-bootloader     filename.ipod\n");
	fprintf(stderr,"  -ab,  --add-bootloader-bin filename.bin\n");
	fprintf(stderr,"  -d,   --delete-bootloader\n");
	fprintf(stderr,"  -f,   --format\n");
	fprintf(stderr,"  -c,   --convert\n");
	fprintf(stderr,"        --read-aupd          filename.bin\n");
	fprintf(stderr,"        --write-aupd         filename.bin\n");
	fprintf(stderr,"\n");
	
#ifdef __WIN32__
	fprintf(stderr,"DISKNO is the number (e.g. 2) Windows has assigned to your iPod's hard disk.\n");
	fprintf(stderr,"The first hard disk in your computer (i.e. C:\\) will be disk 0, the next disk\n");
	fprintf(stderr,"will be disk 1 etc.\n");
#elif defined(linux) || defined (__linux)
	fprintf(stderr,"\"device\" is the device node (e.g. /dev/sda) assigned to your iPod.\n");
#elif defined(__FreeBSD__) || defined(__NetBSD__) || defined(__OpenBSD__)
	fprintf(stderr,"\"device\" is the device node (e.g. /dev/da1) assigned to your iPod.\n");
#elif defined(__APPLE__) && defined(__MACH__)
	fprintf(stderr,"\"device\" is the device node (e.g. /dev/disk1) assigned to your iPod.\n");
#endif
	fprintf(stderr,"ipodpatcher will refuse to access a disk unless it can identify it as being\n");
	fprintf(stderr,"an iPod. ipodpatcher will return -1 if no iPod is found or cannot be \n");
	fprintf(stderr,"accessed, 0 if the requested action is successful, or 1 if unsuccessful.\n");
	fprintf(stderr,"\n");
}

void display_partinfo(struct ipod_t* ipod)
{
	int i;
	double sectors_per_MB = (1024.0*1024.0)/ipod->sector_size;
	
	printf("[INFO] Part    Start Sector    End Sector   Size (MB)   Type\n");
	for ( i = 0; i < 4; i++ ) {
		if (ipod->pinfo[i].start != 0) {
			printf("[INFO]    %d      %10ld    %10ld  %10.1f   %s (0x%02x)\n",
					i,
					(long int)ipod->pinfo[i].start,
					(long int)ipod->pinfo[i].start+ipod->pinfo[i].size-1,
					ipod->pinfo[i].size/sectors_per_MB,
					get_parttype(ipod->pinfo[i].type),
					(int)ipod->pinfo[i].type
					);
        }
    }
}

void exit_prompt()
{
#ifdef WITH_BOOTOBJS
	char enter[4];
	printf("\nPress ENTER to exit ipodpatcher :");
	fgets(enter,4,stdin);
#endif
}

int main(int argc, char* argv[])
{
	char yesno[4];
	int i;
	int n;
	int infile, outfile;
	unsigned int inputsize;
	char* filename;
	int action = NONE;
	int type;
	struct ipod_t ipod;
	
	printf("ipodpatcher v" VERSION " - (C) Dave Chapman 2006-2007\n");
	printf("This is free software; see the source for copying conditions.  There is NO\n");
	printf("warranty; not even for MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.\n\n");
	
	if ((argc > 1) && ((strcmp(argv[1],"-h")==0) || (strcmp(argv[1],"--help")==0))) {
		print_usage();
		return 0;
	}
	
	if (ipod_alloc_buffer(&sectorbuf,BUFFER_SIZE) < 0) {
		fprintf(stderr,"Failed to allocate memory buffer\n");
		return 1;
    }
	
	if ((argc > 1) && (strcmp(argv[1],"--scan")==0)) {
		if (ipod_scan(&ipod) == 0) {
			fprintf(stderr,"[ERR]  No iPod found.\n");
			return -1;
		} else {
			return 0;
		}
	}
	
	/* If the first parameter doesn't start with -, then we interpret it as a device */
    if ((argc > 1) && (argv[1][0] != '-')) {
		ipod.diskname[0]=0;
#ifdef __WIN32__
		snprintf(ipod.diskname,sizeof(ipod.diskname),"\\\\.\\PhysicalDrive%s",argv[1]);
#else
		strncpy(ipod.diskname,argv[1],sizeof(ipod.diskname));
#endif
		i = 2;
	} else {
		/* Autoscan for ipods */
		n = ipod_scan(&ipod);
		if (n==0) {
			fprintf(stderr,"[ERR]  No iPod found: aborting!\n");
			fprintf(stderr,"[ERR]  Please connect your iPod and ensure it is in disk mode\n");
#if defined(__APPLE__) && defined(__MACH__)
			fprintf(stderr,"[ERR]  Also ensure that iTunes is closed, and that your iPod is not mounted.\n");
#elif defined(__WIN32__)
			fprintf(stderr,"[ERR]  You may also need to run ipodpatcher with admin access.\n");
#else
			fprintf(stderr,"[ERR]  You may also need to run ipodpatcher as root.\n");
#endif
			fprintf(stderr,"[ERR]  Please refer to the Rockbox manual if you continue to have problems.\n");
			exit_prompt();
			return -1;
		} else if (n > 1) {
			fprintf(stderr,"[ERR]  %d iPods found: aborting!\n",n);
			fprintf(stderr,"[ERR]  Please connect only one iPod and re-run ipodpatcher.\n");
			exit_prompt();
			return -1;
		}
		
		i = 1;
	}
	
#ifdef WITH_BOOTOBJS
	action = INTERACTIVE;
#else
	action = NONE;
#endif
	
	/* Parse args*/
	while (i < argc) {
		char *arg = argv[i];
		if ((strcmp(arg,"-l")==0) || (strcmp(arg,"--list")==0)) {
			action = LIST_IMAGES;
#ifdef WITH_BOOTOBJS
		} else if (strcmp(argv[i],"--install")==0) {
			action = INSTALL;
#endif
		} else if ((strcmp(arg,"-d")==0) ||
					(strcmp(arg,"--delete-bootloader")==0)) {
			action = DELETE_BOOTLOADER;
		} else if ((strcmp(arg,"-a")==0) || 
					(strcmp(arg,"--add-bootloader")==0)) {
			action = ADD_BOOTLOADER;
			type = FILETYPE_DOT_IPOD;
			i++;
			if (i == argc) {
				print_usage();
				return 1;
				}
			filename=argv[i];
		} else if ((strcmp(arg,"-ab")==0) || 
					(strcmp(arg,"--add-bootloader-bin")==0)) {
			action = ADD_BOOTLOADER;
			type = FILETYPE_DOT_BIN;
			i++;
			if (i == argc) {
				print_usage();
				return 1;
			}
			filename=argv[i];
		} else if ((strcmp(arg,"-rf")==0) || 
					(strcmp(arg,"--read-firmware")==0)) {
			action = READ_FIRMWARE;
			type = FILETYPE_DOT_IPOD;
			i++;
			if (i == argc) {
				print_usage();
				return 1;
			}
			filename=argv[i];
		} else if ((strcmp(arg,"-rfb")==0) || 
					(strcmp(arg,"--read-firmware-bin")==0)) {
			action = READ_FIRMWARE;
			type = FILETYPE_DOT_BIN;
			i++;
			if (i == argc) {
				print_usage();
				return 1;
			}
			filename=argv[i];
#ifdef WITH_BOOTOBJS
		} else if ((strcmp(arg,"-we")==0) || 
					(strcmp(arg,"--write-embedded")==0)) {
			action = WRITE_FIRMWARE;
			type = FILETYPE_INTERNAL;
			filename="[embedded bootloader]";  /* Only displayed for user */
#endif
		} else if ((strcmp(arg,"-wf")==0) || 
					(strcmp(arg,"--write-firmware")==0)) {
			action = WRITE_FIRMWARE;
			type = FILETYPE_DOT_IPOD;
			i++;
			if (i == argc) {
				print_usage();
				return 1;
			}
			filename=argv[i];
		} else if ((strcmp(arg,"-wfb")==0) || 
					(strcmp(arg,"--write-firmware-bin")==0)) {
			action = WRITE_FIRMWARE;
			type = FILETYPE_DOT_BIN;
			i++;
			if (i == argc) {
				print_usage();
				return 1;
				}
			filename=argv[i];
		} else if ((strcmp(arg,"-r")==0) || 
					(strcmp(arg,"--read-partition")==0)) {
			action = READ_PARTITION;
			i++;
			if (i == argc) {
				print_usage();
				return 1;
			}
			filename=argv[i];
		} else if ((strcmp(arg,"-w")==0) || 
					(strcmp(arg,"--write-partition")==0)) {
			action = WRITE_PARTITION;
			i++;
			if (i == argc) {
				print_usage();
				return 1;
			}
			filename=argv[i];
		} else if ((strcmp(arg,"-v")==0) || 
					(strcmp(arg,"--verbose")==0)) {
			verbose=1;
		} else if ((strcmp(arg,"-f")==0) || 
                   (strcmp(arg,"--format")==0)) {
			action = FORMAT_PARTITION;
		} else if (strcmp(arg,"--read-aupd")==0) {
			action = READ_AUPD;
			i++;
			if (i == argc) {
				print_usage();
				return 1;
				}
			filename=argv[i];
        } else if (strcmp(arg,"--write-aupd")==0) {
			action = WRITE_AUPD;
			i++;
			if (i == argc) {
				print_usage();
				return 1;
				}
			filename=argv[i];
		} else if ((strcmp(arg,"-c")==0) || 
					(strcmp(arg,"--convert")==0)) {
			action = CONVERT_TO_FAT32;
		} else {
			print_usage();
			return 1;
		}
		i++;
	}
	
	if (ipod.diskname[0]==0) {
		print_usage();
		return 1;
	}
	
	if (ipod_open(&ipod, 0) < 0) {
		return -1;
	}
	
	fprintf(stderr,"[INFO] Reading partition table from %s\n",ipod.diskname);
	fprintf(stderr,"[INFO] Sector size is %d bytes\n",ipod.sector_size);
	
	if (read_partinfo(&ipod,0) < 0) {
		return -1;
    }
	
	display_partinfo(&ipod);
	
	if (ipod.pinfo[0].start==0) {
		fprintf(stderr,"[ERR]  No partition 0 on disk:\n");
		return -1;
	}
	
	read_directory(&ipod);
	
	if (ipod.nimages <= 0) {
		fprintf(stderr,"[ERR]  Failed to read firmware directory - nimages=%d\n",ipod.nimages);
		return -1;
	}
	
	if (getmodel(&ipod,(ipod.ipod_directory[0].vers>>8)) < 0) {
		fprintf(stderr,"[ERR] Unknown version number in firmware (%08x)\n", ipod.ipod_directory[0].vers);
		return 1;
	}
	
	printf("[INFO] iPod model: %s (\"%s\")\n",ipod.modelstr, ipod.macpod ? "Macpod" : "Winpod");
	
	if (ipod.ipod_directory[0].vers == 0x10000) {
		fprintf(stderr,"[ERR]  *** ipodpatcher does not support the 2nd Generation Nano! ***\n");
		exit_prompt();
		return 1;
	}
	
	if (ipod.macpod) {
		print_macpod_warning();
	}
	
	/* Execute the action */
	if (action > LIST_IMAGES) {
		if (ipod_reopen_rw(&ipod) < 0) {
			return -1;
		}
	}
	switch (action) {
		case LIST_IMAGES:
			list_images(&ipod);
			break;
#ifdef WITH_BOOTOBJS
		case INTERACTIVE:
			printf("Enter i to install the Rockbox bootloader, u to uninstall\n or "
					"c to cancel and do nothing (i/u/c) :");
			if (fgets(yesno,4,stdin)) {
				if (yesno[0]=='i') {
					if (add_bootloader(&ipod, NULL, FILETYPE_INTERNAL)==0) {
						fprintf(stderr,"[INFO] Bootloader installed successfully.\n");
					} else {
						fprintf(stderr,"[ERR]  --install failed.\n");
						return 1;
					}
				} else if (yesno[0]=='u') {
					if (delete_bootloader(&ipod)==0) {
						fprintf(stderr,"[INFO] Bootloader removed.\n");
					} else {
						fprintf(stderr,"[ERR]  Bootloader removal failed.\n");
						return 1;
					}
				} else {
					fprintf(stderr,"[INFO] Installation cancelled.\n");
					return 1;
				}
			}
			exit_prompt();
			break;
		
		case INSTALL:
			if (add_bootloader(&ipod, NULL, FILETYPE_INTERNAL)==0) {
				fprintf(stderr,"[INFO] Bootloader installed successfully.\n");
			} else {
				fprintf(stderr,"[ERR]  --install failed.\n");
				return 1;
			}
			break;
#endif
		case DELETE_BOOTLOADER:
			if (ipod.ipod_directory[0].entryOffset==0) {
				fprintf(stderr,"[ERR]  No bootloader detected.\n");
				return 1;
			} else {
				if (delete_bootloader(&ipod)==0) {
					fprintf(stderr,"[INFO] Bootloader removed.\n");
				} else {
					fprintf(stderr,"[ERR]  --delete-bootloader failed.\n");
					return 1;
				}
			}
			break;
		
		case ADD_BOOTLOADER:
			if (add_bootloader(&ipod, filename, type)==0) {
				fprintf(stderr,"[INFO] Bootloader %s written to device.\n",filename);
			} else {
				fprintf(stderr,"[ERR]  --add-bootloader failed.\n");
				return 1;
			}
			break;
		
		case READ_FIRMWARE:
			if (read_firmware(&ipod, filename, type)==0) {
				fprintf(stderr,"[INFO] Firmware read to file %s.\n",filename);
			} else {
				fprintf(stderr,"[ERR]  --read-firmware failed.\n");
				return 1;
			}
			break;
		
		case WRITE_FIRMWARE:
			if (write_firmware(&ipod, filename,type)==0) {
				fprintf(stderr,"[INFO] Firmware %s written to device.\n",filename);
			} else {
				fprintf(stderr,"[ERR]  --write-firmware failed.\n");
				return 1;
			}
			break;
		
		case READ_AUPD:
			if (read_aupd(&ipod, filename)==0) {
				fprintf(stderr,"[INFO] AUPD image read to file %s.\n",filename);
			} else {
				fprintf(stderr,"[ERR]  --read-aupd failed.\n");
				return 1;
			}
			break;
		
		case WRITE_AUPD:
			if (write_aupd(&ipod, filename)==0) {
				fprintf(stderr,"[INFO] AUPD image %s written to device.\n",filename);
			} else {
				fprintf(stderr,"[ERR]  --write-aupd failed.\n");
				return 1;
			}
			break;
		
		case READ_PARTITION:
			outfile = open(filename,O_CREAT|O_TRUNC|O_WRONLY|O_BINARY,S_IREAD|S_IWRITE);
			if (outfile < 0) {
				perror(filename);
				return 1;
			}
			if (read_partition(&ipod, outfile)==0) {
				fprintf(stderr,"[INFO] Partition extracted to %s.\n",filename);
			} else {
				fprintf(stderr,"[ERR]  --read-partition failed.\n");
				return 1;
			}
			close(outfile);
			break;
		
		case WRITE_PARTITION:
			infile = open(filename,O_RDONLY|O_BINARY);
			if (infile < 0) {
				perror(filename);
				return 1;
			}
			/* Check filesize is <= partition size */
			inputsize=filesize(infile);
			if (inputsize > 0) {
				if (inputsize <= (ipod.pinfo[0].size*ipod.sector_size)) {
					fprintf(stderr,"[INFO] Input file is %u bytes\n",inputsize);
					if (write_partition(&ipod,infile)==0) {
						fprintf(stderr,"[INFO] %s restored to partition\n",filename);
					} else {
						fprintf(stderr,"[ERR]  --write-partition failed.\n");
						return 1;
					}
				} else {
					fprintf(stderr,"[ERR]  File is too large for firmware partition: aborting!\n");
					return 1;
				}
			}
			close(infile);
			break;
		
		case FORMAT_PARTITION:
			printf("WARNING!!! YOU ARE ABOUT TO USE AN EXPERIMENTAL FEATURE.\n");
			printf("ALL DATA ON YOUR IPOD WILL BE ERASED.\n");
			printf("Are you sure you want to format your iPod? (y/n):");
			if (fgets(yesno,4,stdin)) {
				if (yesno[0]=='y') {
					if (format_partition(&ipod,1) < 0) {
						fprintf(stderr,"[ERR]  Format failed.\n");
						return 1;
					}
				} else {
					fprintf(stderr,"[INFO] Format cancelled.\n");
					return 1;
				}
			}
			break;
		
		case CONVERT_TO_FAT32:
			if (!ipod.macpod) {
				fprintf(stderr,"[ERR]  iPod is already FAT32: aborting!\n");
			} else {
				printf("WARNING!!! YOU ARE ABOUT TO USE AN EXPERIMENTAL FEATURE.\n");
				printf("ALL DATA ON YOUR IPOD WILL BE ERASED.\n");
				printf("Are you sure you want to convert your iPod to FAT32? (y/n):");
				if (fgets(yesno,4,stdin)) {
					if (yesno[0]=='y') {
						if (write_dos_partition_table(&ipod) < 0) {
							fprintf(stderr,"[ERR]  Partition conversion failed.\n");
							return 1;
						}
						if (format_partition(&ipod,1) < 0) {
							fprintf(stderr,"[ERR]  Format failed.\n");
							return 1;
						}
					} else {
						fprintf(stderr,"[INFO] Format cancelled.\n");
						return 1;
					}
				}
			}
			break;
		default:
			break;
	}
	
	ipod_close(&ipod);
	return 0; /* If nothing went wrong, assume success */
	
}