// Podtronics Software License
// Version 1.0
// Copyright (c) 2002 Grant Jones. All rights reserved.
// 
// 1. Redistributions of source code must retain the above copyright notice, this
// list of conditions and the following disclaimer.
// 
// 2. Altered source versions must be plainly marked as such, and must not be
// misrepresented as being the original software.
// 
// 3. 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.
// 
// 4. The end-user documentation included with the redistribution, if any, must
// include the following acknowledgment:
// 
// "This product includes software developed by Grant Jones (http://www.podtronics.com/)."
// 
// Alternately, this acknowledgment may appear in the software itself, if and
// wherever such third-party acknowledgments normally appear.
// 
// 5. This source code may be used free of charge in an end product that is free of
// cost. Where the end product is not free of cost, a commercial license fee must be
// paid to the original developer.
// 
// THIS SOFTWARE IS PROVIDED "AS IS" AND ANY EXPRESSED 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 
// GRANT JONES, PODTRONICS, OR CONTRIBUTORS 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.
// 
// This license is based upon the Apache Software License Version 1.1
// (http://www.apache.org/).
// 
// ------------------------------------------------------------------------------
// 
// Updates
// Date     Author        Comment
// ------------------------------------------------------------------------------
// 1/13/02  Grant         Original Cut
// 1/13/02  Grant         Ported to Linux (May need a little more hacking to get it to work)


// ////////////////////////////////////////////////////////////////
// iPodUpdater.cpp : Defines the entry point for the iPod Updater.


// C RunTime Header Files
#include <stdlib.h>
#include <stdio.h>
#include <malloc.h>


#define DOING_CHECKRUN		// Un-comment to disable writing to the iPod

// SOme lazy typedefitions
typedef int HANDLE;
typedef unsigned char BOOL;
typedef unsigned long DWORD;

/// Firmware upgrade modes
enum
{
	MODE_INVALID = 0,
	MODE_UPDATE = 10,
	MODE_DUMP = 12,
	MODE_PROBE = 14
};
//////////////////////////////////////////////////////////////////////////////
// Partition Map Info
//////////////////////////////////////////////////////////////////////////////
typedef struct
{
	unsigned short	pmSig;						// {partition signature}
	unsigned short	pmSigPad;					// {reserved}
	unsigned long	pmMapBlkCnt;					// {number of blocks in partition map}
	unsigned long	pmPyPartStart;				// {first physical block of partition}
	unsigned long	pmPartBlkCnt;					// {number of blocks in partition}
    unsigned char	pmPartName[32];				// {partition name}
	unsigned char	pmParType[32];				// {partition type}
	unsigned long	pmLgDataStart;				// {first logical block of data area}
	unsigned long	pmDataCnt;					// {number of blocks in data area}
	unsigned long	pmPartStatus;					// {partition status information}
	unsigned long	pmLgBootStart;				// {first logical block of boot code}
	unsigned long	pmBootSize;					// {size of boot code, in bytes}
	unsigned long	pmBootAddr;					// {boot code load address}
	unsigned long	pmBootAddr2;					// {reserved}
	unsigned long	pmBootEntry;					// {boot code entry point}
	unsigned long	pmBootEntry2;					// {reserved}
	unsigned long	pmBootCksum;					// {boot code checksum}
	unsigned char	pmProcessor[16];				// {processor type}
	unsigned short	pmPad[188];					// {reserved}
} PartitionMapInfo;

// Byte swapping stuff
#define swap2long(x) x =	( (x >> 24) & 0x000000ffL ) +			\
							( (x >> 8) & 0x0000ff00L) +				\
							( (x << 8) & 0x00ff0000L) +				\
							( (x << 24) & 0xff000000L);

#define swap2short(x) x =	(((x >> 8) & 0x00ff) + ((x << 8) & 0xff00));

unsigned char llread( HANDLE disk, unsigned char * bufPtr, DWORD offset, DWORD length );
BOOL llwrite( HANDLE disk, unsigned char *bufPtr, DWORD offset, DWORD length );
BOOL dprintf( char *fmt, ... );
void swappartitionmap( PartitionMapInfo *pmi );

////////////////////////////////////////////////////////////////////////////////////////

int main( int argc, char *argv[] )
{
	char firmware[200];		// Firmware file
	char path[200];			// Path to iPod drive

	dprintf("Create by Grant Jones (grant@podtronics.com)");
	dprintf("iPod Updater v1.1 Beta");
	dprintf("Go to podtronics.com for updates!");
	dprintf("------------------------------------------------");
	dprintf("Backup your partition before doing any thing and remember:");
	dprintf("I am not responsible for the loss of data caused by this program!");

	updateFirmware( firmware, path, MODE_UPDATE );

	return 0;
}

BOOL dprintf( char *fmt, ... )
{
	va_list argptr; // Pointer to argument list
	char str[ 140 ];
	int cnt; // Result of sprintf

	va_start( argptr, fmt );
	cnt = vsprintf( str, fmt, argptr );
	if( str[0] == NULL )
		return FALSE;
	
	// Add text here
	if( str[strlen(str)-1] == '\n' )
		str[strlen(str)-1] = '\0';
	
	fprintf(stdout, "%s\n", str );			// print string
	
	va_end( argptr );
	return TRUE;

}


void swappartitionmap( PartitionMapInfo *pmi )
{
	pmi->pmSig = swap2short( pmi->pmSig );							// {partition signature}
	pmi->pmSigPad = swap2short( pmi->pmSigPad );					// {reserved}
	pmi->pmMapBlkCnt = swap2long( pmi->pmMapBlkCnt );				// {number of blocks in partition map}
	pmi->pmPyPartStart = swap2long( pmi->pmPyPartStart );			// {first physical block of partition}
	pmi->pmPartBlkCnt = swap2long( pmi->pmPartBlkCnt );				// {number of blocks in partition}
    pmi->pmLgDataStart = swap2long( pmi->pmLgDataStart );			// {first logical block of data area}
	pmi->pmDataCnt = swap2long( pmi->pmDataCnt );					// {number of blocks in data area}
	pmi->pmPartStatus = swap2long( pmi->pmPartStatus );				// {partition status information}
	pmi->pmLgBootStart = swap2long( pmi->pmLgBootStart );			// {first logical block of boot code}
	pmi->pmBootSize = swap2long( pmi->pmBootSize );					// {size of boot code, in bytes}
	pmi->pmBootAddr = swap2long( pmi->pmBootAddr );					// {boot code load address}
	pmi->pmBootAddr2 = swap2long( pmi->pmBootAddr2 );				// {reserved}
	pmi->pmBootEntry = swap2long( pmi->pmBootEntry );				// {boot code entry point}
	pmi->pmBootEntry2 = swap2long( pmi->pmBootEntry2 );				// {reserved}
	pmi->pmBootCksum = swap2long( pmi->pmBootCksum );				// {boot code checksum}
}


#ifdef STILL_NEEDS_PORTING
char findiPod( void )
{
	HANDLE hDrive;
	char createFileString[] = "\\\\.\\PHYSICALDRIVE0";
	DWORD offset;
	PartitionMapInfo pmi;
	char i;
	
	// Scan 27 physical drives trying to find the iPod
	for( i = 0; i < 27; i ++ )
	{
		sprintf( createFileString, "\\\\.\\PHYSICALDRIVE%d", i );
		
		hDrive = CreateFile( createFileString, GENERIC_READ|GENERIC_WRITE, 0, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL | FILE_FLAG_WRITE_THROUGH | FILE_FLAG_NO_BUFFERING, NULL );
		if( INVALID_HANDLE_VALUE == hDrive )
		{
			continue;
		}
		
		// Find the partition map entry and load the values for that
		offset = 512;

		// Need to find where the correct partition is
		if( llread( hDrive, (unsigned char *)&pmi, offset, 512 ) != 255 )
		{
			continue;
		}
	
		swappartitionmap( &pmi );
		if( strcmp( (char *)pmi.pmPartName, "partition map" ) == 0 )
		{
			dprintf("Found iPod Drive!");
			CloseHandle( hDrive );
			return i;
		}
		else
		{
			CloseHandle( hDrive );
			continue;
		}
	}
	return -1;
}
#endif

BOOL updateFirmware( char *firmware, char *path, char mode )
{
	HANDLE hDrive;
	//char createFileString[] = "\\\\.\\PHYSICALDRIVE0";
	DWORD offset;
	//char PhysicaliPodDrive = 0;
	PartitionMapInfo pmi;
	PartitionMapInfo *pmis;
	FILE *fp;
	FILE *foo;
	int y;
	int red;
	unsigned char buf[512];
	BOOL success = FALSE;
	int curpart = 0;
	
	dprintf(" *****************************************");

	if( mode == MODE_INVALID )
		return FALSE;
#ifdef FOOBAR
	PhysicaliPodDrive = findiPod();
	if( PhysicaliPodDrive < 0 )
	{
		dprintf("Unable to find the iPod! Make sure it is connect and powered!");
		return FALSE;
	}
	
	sprintf( createFileString, "\\\\.\\PHYSICALDRIVE%d", PhysicaliPodDrive );

	hDrive = CreateFile( createFileString, GENERIC_READ|GENERIC_WRITE, 0, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL | FILE_FLAG_WRITE_THROUGH | FILE_FLAG_NO_BUFFERING, NULL );
	if( INVALID_HANDLE_VALUE == hDrive )
	{
		dprintf("Unable to open physical drive %d (gle: %ld)", PhysicaliPodDrive, GetLastError() );
		if( mode != MODE_PROBE )
			MessageBox( NULL, "Unable to open iPod drive!!", "Error", MB_OK );
		return FALSE;
	}
#endif
	hDrive = open( path, O_RDWR, 0 );
	if( hDrive <= 0 )
	{
		dprintf("Couldn't open drive!");
		return FALSE;
	}

	dprintf("loading partitions...");

	// Find the partition map entry and load the values for that
	offset = 512;

	// Need to find where the correct partition is
	if( llread( hDrive, (unsigned char *)&pmi, offset, 512 ) != 255 )
	{
		dprintf("Drive read failed! Dc:1");
		close(hDrive);
		return FALSE;
	}
	
	swappartitionmap( &pmi );
	if( strcmp( (char *)pmi.pmPartName, "partition map" ) == 0 )
	{
		//dprintf("Drive %c is an apple drive!", iPodDrive);
	}
	else
	{
		dprintf("Physical Drive %d is not an apple drive!", PhysicaliPodDrive);
		dprintf("\"%s\" part found.", pmi.pmPartName );
		close(hDrive);
		return FALSE;
	}

	pmis = (PartitionMapInfo *)malloc( sizeof( PartitionMapInfo ) * (pmi.pmMapBlkCnt) );
	if(!pmis)
		return FALSE;

	memcpy( &pmis[0], &pmi, sizeof( PartitionMapInfo ) );

	for( unsigned int i = 1; i < pmi.pmMapBlkCnt; i ++ )
	{
		offset += 512;
		
		// Need to find where the correct partition is
		if( llread( hDrive, (unsigned char *)&pmis[i], offset, 512 ) != 255 )
		{
			dprintf("Drive read failed! Dc:2");
			break;
		}
	
		swappartitionmap( &pmis[i] );
		
		if( strcmp( (char *)pmis[i].pmPartName, "firmware" ) == 0 )
		{
			dprintf("Found firmware on iPod!");

			dprintf("The size of the partition is %ld bytes.", pmis[i].pmPartBlkCnt*512 );
			if( mode == MODE_PROBE )
			{
				success = TRUE;
				break;
			}

			if( mode == MODE_DUMP )
			{
				unsigned int x;
				BOOL inc = FALSE;
				FILE *dump;
				dprintf("Saving firmware to \"%s\"...", firmware );
				dump = fopen( firmware, "wb" );
				if(!dump)
				{
					dprintf("Error: Unable to open firmware file!");
					break;
				}
				setpbrange( 0, pmis[i].pmPartBlkCnt );
				offset = pmis[i].pmPyPartStart*512;

				for( x = 0; x < pmis[i].pmPartBlkCnt; x++ )
				{
					
					llread( hDrive, buf, offset+(x*512), 512 );
					fwrite( buf, sizeof(unsigned char), 512, dump );

					incrementProgressbar();
				}

				fclose(dump);
				dprintf("Finished! File saved to \"%s\"", firmware );
				success = TRUE;
				break;
			}

			// If we're here we are updating the firmware

			dprintf("Updating firmware with file \"%s\"...", firmware );			
			
			fp = fopen( firmware, "rb" );
			if(!fp)
			{
				dprintf("Unable to open firmware file!");
				break;
			}
#ifdef DOING_CHECKRUN
			foo = fopen( "checkrun", "wb" );
#endif
			offset = pmis[i].pmPyPartStart*512;
			y = 0;
			success = TRUE;

			while(!feof(fp))
			{
				red = fread( buf, 1, 512, fp );
#ifndef DOING_CHECKRUN
				if( !llwrite( hDrive, buf, offset+(y*512), red ) )
				{
					success = FALSE;
					dprintf("Error quiting");
					break;
				}
#else
				fwrite( buf, 1, red, foo );
#endif
				
				y++;
			}

			fclose(fp);
#ifdef DOING_CHECKRUN
			fclose(foo);
#endif
			if( success == TRUE )
				dprintf("Updated firmware!");
			else
				dprintf("Unable to update firmware!");
			
			break;
		}
				
	}

	
	
	free( pmis );
	close(hDrive);
	return success;
}


// These read and writes have no handling for block aligned stuff
// If you need a read and write that is block aligned see Apples darwin code :)

unsigned char llread( HANDLE disk, unsigned char *bufPtr, DWORD offset, DWORD length )
{	
	DWORD		blocksize;
	DWORD		readResult;

	blocksize = 512;	// This is correct for an iPod drive

	if( (length % blocksize)!=0 ||
		(offset % blocksize)!= 0 )
	{
		return 1;
	}

	readResult = lseek( disk, offset, SEEK_SET );
    if ( readResult != offset )
	{
		return 2;
	}
	
	readResult = read(disk, bufPtr, length);
	if( readResult != length )
	{
		return 3;
	}
	return 255;
}

BOOL llwrite( HANDLE disk, unsigned char *bufPtr, DWORD offset, DWORD length )
{
	DWORD		blocksize;
	DWORD		writeResult;

	blocksize = 512;	// This is correct for an iPod drive

	if( (length % blocksize)!=0 ||
		(offset % blocksize)!= 0 )
	{
		return FALSE;
	}

	writeResult = lseek( disk, offset, SEEK_SET );
    if ( writeResult != offset )
	{
		return FALSE;
	}
	
	writeResult = write(disk, bufPtr, length);
	if( writeResult != length )
	{
		return FALSE;
	}
	return FALSE;
}