Start to implement some functionality. You can now read/write the entire partition (as ipodpatcher used to do), extract the original firmware to an "apple_os.ipod" file, or replace the original firmware with either a Rockbox bootloader or Rockbox itself. Adding/Removing a bootloader isn't yet implemented.

git-svn-id: svn://svn.rockbox.org/rockbox/trunk@11755 a1c6a512-1295-4272-9138-f99709370657

1 files changed, 418 insertions, 66 deletions

diff --git a/tools/ipodpatcher/ipodpatcher.c b/tools/ipodpatcher/ipodpatcher.c
index f965746282..bfad75b98a 100644
--- a/tools/ipodpatcher/ipodpatcher.c
+++ b/tools/ipodpatcher/ipodpatcher.c
@@ -59,8 +59,9 @@ static const char *apple_stop_sign = "{{~~  /-----\\   "\
    and initialise it with ipod_alloc_buf() 
 */
 
-/* Size of buffer for disk I/O */
-#define BUFFER_SIZE 6*1024*1024
+/* Size of buffer for disk I/O - 8MB is large enough for any version
+   of the Apple firmware, but not the Nano's RSRC image. */
+#define BUFFER_SIZE 8*1024*1024
 unsigned char* sectorbuf;
 
 char* get_parttype(int pt)
@@ -101,13 +102,20 @@ struct partinfo_t {
   unsigned char type;
 };
 
-int static inline getint32le(unsigned char* buf)
+int static inline le2int(unsigned char* buf)
 {
    int32_t res = (buf[3] << 24) | (buf[2] << 16) | (buf[1] << 8) | buf[0];
 
    return res;
 }
 
+int static inline be2int(unsigned char* buf)
+{
+   int32_t res = (buf[0] << 24) | (buf[1] << 16) | (buf[2] << 8) | buf[3];
+
+   return res;
+}
+
 int static inline getint16le(char* buf)
 {
    int16_t res = (buf[1] << 8) | buf[0];
@@ -115,6 +123,27 @@ int static inline getint16le(char* buf)
    return res;
 }
 
+void static inline short2le(unsigned short val, unsigned char* addr)
+{
+    addr[0] = val & 0xFF;
+    addr[1] = (val >> 8) & 0xff;
+}
+
+void static inline int2le(unsigned int val, unsigned char* addr)
+{
+    addr[0] = val & 0xFF;
+    addr[1] = (val >> 8) & 0xff;
+    addr[2] = (val >> 16) & 0xff;
+    addr[3] = (val >> 24) & 0xff;
+}
+
+void int2be(unsigned int val, unsigned char* addr)
+{
+    addr[0] = (val >> 24) & 0xff;
+    addr[1] = (val >> 16) & 0xff;
+    addr[2] = (val >> 8) & 0xff;
+    addr[3] = val & 0xFF;
+}
 
 
 #define BYTES2INT32(array,pos)\
@@ -157,13 +186,13 @@ int read_partinfo(HANDLE dh, int sector_size, struct partinfo_t* pinfo)
     /* check that the boot sector is initialized */
     if ( (sector[510] != 0x55) ||
          (sector[511] != 0xaa)) {
-        fprintf(stderr,"Bad boot sector signature\n");
-        return 0;
+        fprintf(stderr,"[ERR]  Bad boot sector signature\n");
+        return -1;
     }
 
     if ((memcmp(&sector[71],"iPod",4) != 0) &&
         (memcmp(&sector[0x40],"This is your Apple iPod. You probably do not want to boot from it!",66) != 0) ) {
-        fprintf(stderr,"Drive is not an iPod, aborting\n");
+        fprintf(stderr,"[ERR]  Drive is not an iPod, aborting\n");
         return -1;
     }
 
@@ -183,8 +212,8 @@ int read_partinfo(HANDLE dh, int sector_size, struct partinfo_t* pinfo)
     return 0;
 }
 
-int disk_read(HANDLE dh, int outfile,unsigned long start, unsigned long count,
-              int sector_size)
+int read_partition(HANDLE dh, int outfile,unsigned long start, 
+                   unsigned long count, int sector_size)
 {
     int res;
     unsigned long n;
@@ -214,7 +243,9 @@ int disk_read(HANDLE dh, int outfile,unsigned long start, unsigned long count,
         }
 
         if (n < chunksize) {
-            fprintf(stderr,"[ERR]  Short read in disk_read() - requested %d, got %lu\n",chunksize,n);
+            fprintf(stderr,
+              "[ERR]  Short read in disk_read() - requested %d, got %lu\n",
+              chunksize,n);
             return -1;
         }
 
@@ -228,7 +259,8 @@ int disk_read(HANDLE dh, int outfile,unsigned long start, unsigned long count,
         }
 
         if (res != n) {
-            fprintf(stderr,"Short write - requested %lu, received %d - aborting.\n",n,res);
+            fprintf(stderr,
+              "Short write - requested %lu, received %d - aborting.\n",n,res);
             return -1;
         }
     }
@@ -237,7 +269,7 @@ int disk_read(HANDLE dh, int outfile,unsigned long start, unsigned long count,
     return 0;
 }
 
-int disk_write(HANDLE dh, int infile,unsigned long start, int sector_size)
+int write_partition(HANDLE dh, int infile,unsigned long start, int sector_size)
 {
     unsigned long res;
     int n;
@@ -281,14 +313,15 @@ int disk_write(HANDLE dh, int infile,unsigned long start, int sector_size)
         }
 
         if (res != n) {
-            fprintf(stderr,"Short write - requested %d, received %lu - aborting.\n",n,res);
+            fprintf(stderr,"[ERR]  Short write - requested %d, received %lu - aborting.\n",n,res);
             return -1;
         }
 
         byteswritten += res;
     }
 
-    fprintf(stderr,"[INFO] Wrote %d bytes plus %d bytes padding.\n",byteswritten-padding,padding);
+    fprintf(stderr,"[INFO] Wrote %d bytes plus %d bytes padding.\n",
+            byteswritten-padding,padding);
     return 0;
 }
 
@@ -300,15 +333,16 @@ void print_usage(void) {
     fprintf(stderr,"Usage: ipodpatcher device [action]\n");
 #endif
     fprintf(stderr,"\n");
-    fprintf(stderr,"Where [action] is one of:\n");
-#if 0
-    fprintf(stderr,"     -e --extract-firmware filename.bin - extract firmware to a file\n");
-    fprintf(stderr,"     -i --insert-firmware  filename.bin - replace the firmware with the file\n");
-    fprintf(stderr,"     -a --add-bootloader   filename.bin - add a bootloader\n");
-    fprintf(stderr,"     -r --remove-bootloader             - remove a bootloader\n");
-#endif
-    fprintf(stderr,"     -l --list                          - list images in firmware partition\n");
+    fprintf(stderr,"Where [action] is one of the following options:\n");
+    fprintf(stderr,"  -l,  --list\n");
+    fprintf(stderr,"  -r,  --read-partition   bootpartition.bin\n");
+    fprintf(stderr,"  -w,  --write-partition  bootpartition.bin\n");
+    fprintf(stderr,"  -ef, --extract-firmware filename.ipod\n");
+    fprintf(stderr,"  -rf, --replace-firmware filename.ipod\n");
+    fprintf(stderr,"  -a,  --add-bootloader   filename.ipod\n");
+    fprintf(stderr,"  -d,  --delete-bootloader\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 disk0, the next disk\n");
@@ -325,7 +359,13 @@ void print_usage(void) {
 enum {
    NONE,
    SHOW_INFO,
-   LIST_IMAGES
+   LIST_IMAGES,
+   REMOVE_BOOTLOADER,
+   INSERT_BOOTLOADER,
+   EXTRACT_FIRMWARE,
+   REPLACE_FIRMWARE,
+   READ_PARTITION,
+   WRITE_PARTITION
 };
 
 char* ftypename[] = { "OSOS", "RSRC", "AUPD", "HIBE" };
@@ -340,7 +380,7 @@ enum firmwaretype_t {
 struct ipod_directory_t {
   enum firmwaretype_t ftype;
   int id;
-  uint32_t devOffset;
+  uint32_t devOffset; /* Offset of image relative to one sector into bootpart*/
   uint32_t len;
   uint32_t addr;
   uint32_t entryOffset;
@@ -349,14 +389,209 @@ struct ipod_directory_t {
   uint32_t loadAddr;
 };
 
-int read_directory(HANDLE dh, int start, int sector_size, struct ipod_directory_t* ipod_directory)
+int remove_bootloader(HANDLE dh, int start, int sector_size, 
+                      struct ipod_directory_t* ipod_directory)
+{
+    fprintf(stderr,"[ERR]  Sorry, not yet implemented.\n");
+    return -1;
+}
+
+int replace_firmware(HANDLE dh, char* filename, int start, int sector_size, 
+                     int nimages, struct ipod_directory_t* ipod_directory, 
+                     off_t diroffset, int modelnum, char* modelname)
+{
+    int length;
+    int i;
+    int n;
+    int infile;
+    int newsize;
+    int bytesavailable;
+    unsigned long chksum=0;
+    unsigned long filechksum=0;
+    unsigned char header[8];  /* Header for .ipod file */
+
+    /* First check that the input file is the correct type for this ipod. */
+    infile=open(filename,O_RDONLY);
+    if (infile < 0) {
+        fprintf(stderr,"[ERR]  Couldn't open input file %s\n",filename);
+    }
+
+    n = read(infile,header,8);
+    if (n < 8) {
+        fprintf(stderr,"[ERR]  Failed to read header from %s\n",filename);
+    }
+
+    if (memcmp(header+4,modelname,4)!=0) {
+        fprintf(stderr,"[ERR]  Model name in input file (%c%c%c%c) doesn't match ipod model (%s)\n",
+                header[4],header[5],header[6],header[7],modelname);
+        close(infile);
+        return -1;
+    }
+
+    filechksum = be2int(header);
+
+    length=filesize(infile)-8;
+    newsize=(length+sector_size-1)&~(sector_size-1);
+
+    fprintf(stderr,"[INFO] Padding input file from 0x%08x to 0x%08x bytes\n",
+            length,newsize);
+
+    if (newsize > BUFFER_SIZE) {
+        fprintf(stderr,"[ERR]  Input file too big for buffer\n");
+        close(infile);
+        return -1;
+    }
+
+    /* Check if we have enough space */
+    /* TODO: Check the size of the partition. */
+    if (nimages > 1) {
+        bytesavailable=ipod_directory[1].devOffset-ipod_directory[0].devOffset;
+        if (bytesavailable < newsize) {
+            fprintf(stderr,"[INFO] Moving images to create room for new firmware...\n");
+
+            /* TODO: Implement image movement */
+            fprintf(stderr,"[ERR]  Image movement not yet implemented.\n");
+            close(infile);
+            return -1;
+        }
+    }
+
+    fprintf(stderr,"[INFO] Reading input file...\n");
+    /* We now know we have enough space, so write it. */
+    memset(sectorbuf+length,0,newsize-length);
+    n = read(infile,sectorbuf,length);
+    if (n < 0) {
+        fprintf(stderr,"[ERR]  Couldn't read input file\n");
+        close(infile);
+        return -1;
+    }
+    close(infile);
+
+    chksum = modelnum;
+    for (i = 0; i < length; i++) {
+         /* add 8 unsigned bits but keep a 32 bit sum */
+         chksum += sectorbuf[i];
+    }
+
+    if (chksum == filechksum) {
+        fprintf(stderr,"[INFO] Checksum OK in %s\n",filename);
+    } else {
+        fprintf(stderr,"[ERR]  Checksum in %s failed check\n",filename);
+        return -1;
+    }
+
+    if (ipod_seek(dh,start+sector_size+ipod_directory[0].devOffset) < 0) {
+        fprintf(stderr,"[ERR]  Seek failed\n");
+        return -1;
+    }
+
+    if ((n = ipod_write(dh,sectorbuf,newsize)) < 0) {
+        perror("[ERR]  Write failed\n");
+        return -1;
+    }
+
+    if (n < newsize) {
+        fprintf(stderr,"[ERR]  Short write - requested %d bytes, received %d\n"
+                      ,i,n);
+        return -1;
+    }
+    fprintf(stderr,"[INFO]  Wrote %d bytes to firmware partition\n",n);
+
+    /* Now we need to update the "len", "entryOffset" and "chksum" fields */
+    chksum = 0;
+    for (i = 0; i < length; i++) {
+         /* add 8 unsigned bits but keep a 32 bit sum */
+         chksum += sectorbuf[i];
+    }
+
+    /* Read directory */
+    if (ipod_seek(dh,start + diroffset) < 0) { return -1; }
+
+    n=ipod_read(dh, sectorbuf, sector_size);
+    if (n < 0) { return -1; }
+
+    /* Update entries for image 0 */
+    int2le(length,sectorbuf+16);
+    int2le(0,sectorbuf+24);
+    int2le(chksum,sectorbuf+28);
+
+    /* Write directory */    
+    if (ipod_seek(dh,start + diroffset) < 0) { return -1; }
+    n=ipod_write(dh, sectorbuf, sector_size);
+    if (n < 0) { return -1; }
+
+    return 0;
+}
+
+int extract_firmware(HANDLE dh, char* filename, int start, int sector_size, 
+                     struct ipod_directory_t* ipod_directory, 
+                     int modelnum, char* modelname)
+{
+    int length;
+    int i;
+    int outfile;
+    int n;
+    unsigned long chksum=0;   /* 32 bit checksum - Rockbox .ipod style*/
+    unsigned char header[8];  /* Header for .ipod file */
+
+    if (ipod_directory[0].entryOffset != 0) {
+        /* We have a bootloader... */
+        length = ipod_directory[0].entryOffset;
+    } else {
+        length = ipod_directory[0].len;
+    }
+
+    fprintf(stderr,"[INFO] Reading firmware (%d bytes)\n",length);
+
+    if (ipod_seek(dh,start+sector_size+ipod_directory[0].devOffset) < 0) {
+        return -1;
+    }
+
+    i = (length+sector_size-1) & ~(sector_size-1);
+    fprintf(stderr,"[INFO] Padding read from 0x%08x to 0x%08x bytes\n",
+            length,i);
+
+    if ((n = ipod_read(dh,sectorbuf,i)) < 0) {
+        return -1;
+    }
+
+    if (n < i) {
+        fprintf(stderr,"[ERR]  Short read - requested %d bytes, received %d\n",
+                i,n);
+        return -1;
+    }
+
+    chksum = modelnum;
+    for (i = 0; i < length; i++) {
+         /* add 8 unsigned bits but keep a 32 bit sum */
+         chksum += sectorbuf[i];
+    }
+
+    int2be(chksum,header);
+    memcpy(header+4,modelname,4);
+
+    outfile = open(filename,O_CREAT|O_WRONLY|O_BINARY,0666);
+    if (outfile < 0) {
+        fprintf(stderr,"[ERR]  Couldn't open file %s\n",filename);
+        return -1;
+    }
+
+    write(outfile,header,8);
+    write(outfile,sectorbuf,ipod_directory[0].len);
+    close(outfile);
+
+    return 0;
+}
+
+int read_directory(HANDLE dh, int start, int sector_size, 
+                   struct ipod_directory_t* ipod_directory, off_t* diroffset)
 {
     int n;
     int nimages;
-    off_t diroffset;
     unsigned char* p;
 
-    /* Read firmware partition header (first 512 bytes of disk - but let's read a whole sector) */
+    /* Read firmware partition header (first 512 bytes of disk - but 
+       let's read a whole sector) */
 
     if (ipod_seek(dh, start) < 0) { return -1; }
 
@@ -369,14 +604,14 @@ int read_directory(HANDLE dh, int start, int sector_size, struct ipod_directory_
     }
 
     if (memcmp(sectorbuf+0x100,"]ih[",4)!=0) {
-        fprintf(stderr,"[ERR] Bad firmware directory\n");
+        fprintf(stderr,"[ERR]  Bad firmware directory\n");
         return -1;
     }
 
-    diroffset=getint32le(sectorbuf+0x104) + 0x200;
+    *diroffset=le2int(sectorbuf+0x104) + 0x200;
 
     /* Read directory */
-    if (ipod_seek(dh,start + diroffset) < 0) { return -1; }
+    if (ipod_seek(dh,start + *diroffset) < 0) { return -1; }
 
     n=ipod_read(dh, sectorbuf, sector_size);
     if (n < 0) { return -1; }
@@ -384,7 +619,8 @@ int read_directory(HANDLE dh, int start, int sector_size, struct ipod_directory_
     nimages=0;
     p = sectorbuf;
     
-    while ((nimages < MAX_IMAGES) && (p < (sectorbuf + 400)) && (memcmp(p,"!ATA",4)==0)) {
+    while ((nimages < MAX_IMAGES) && (p < (sectorbuf + 400)) && 
+           (memcmp(p,"!ATA",4)==0)) {
         p+=4;
         if (memcmp(p,"soso",4)==0) {
             ipod_directory[nimages].ftype=FTYPE_OSOS;
@@ -395,38 +631,40 @@ int read_directory(HANDLE dh, int start, int sector_size, struct ipod_directory_
         } else if (memcmp(p,"ebih",4)==0) {
             ipod_directory[nimages].ftype=FTYPE_HIBE;
         } else {
-            fprintf(stderr,"[ERR]  Unknown image type %c%c%c%c\n",p[0],p[1],p[2],p[3]);
+            fprintf(stderr,"[ERR]  Unknown image type %c%c%c%c\n",
+                           p[0],p[1],p[2],p[3]);
         }
         p+=4;
-        ipod_directory[nimages].id=getint32le(p);
+        ipod_directory[nimages].id=le2int(p);
         p+=4;
-        ipod_directory[nimages].devOffset=getint32le(p);
+        ipod_directory[nimages].devOffset=le2int(p);
         p+=4;
-        ipod_directory[nimages].len=getint32le(p);
+        ipod_directory[nimages].len=le2int(p);
         p+=4;
-        ipod_directory[nimages].addr=getint32le(p);
+        ipod_directory[nimages].addr=le2int(p);
         p+=4;
-        ipod_directory[nimages].entryOffset=getint32le(p);
+        ipod_directory[nimages].entryOffset=le2int(p);
         p+=4;
-        ipod_directory[nimages].chksum=getint32le(p);
+        ipod_directory[nimages].chksum=le2int(p);
         p+=4;
-        ipod_directory[nimages].vers=getint32le(p);
+        ipod_directory[nimages].vers=le2int(p);
         p+=4;
-        ipod_directory[nimages].loadAddr=getint32le(p);
+        ipod_directory[nimages].loadAddr=le2int(p);
         p+=4;
         nimages++;
     }
     return nimages;
 }
 
-int list_images(int nimages, struct ipod_directory_t* ipod_directory)
+int list_images(int nimages, struct ipod_directory_t* ipod_directory, 
+                int sector_size)
 {
     int i;
 
 #ifdef DEBUG
-    printf("    Type         id  devOffset        len       addr entryOffset    chksum       vers   loadAddr\n");
+    printf("    Type         id  devOffset        len       addr entryOffset    chksum       vers   loadAddr   devOffset+len\n");
     for (i = 0 ; i < nimages; i++) {
-        printf("%d - %s 0x%08x 0x%08x 0x%08x 0x%08x 0x%08x 0x%08x 0x%08x 0x%08x\n",i,
+        printf("%d - %s 0x%08x 0x%08x 0x%08x 0x%08x 0x%08x 0x%08x 0x%08x 0x%08x 0x%08x\n",i,
                ftypename[ipod_directory[i].ftype],
                ipod_directory[i].id,
                ipod_directory[i].devOffset,
@@ -435,7 +673,8 @@ int list_images(int nimages, struct ipod_directory_t* ipod_directory)
                ipod_directory[i].entryOffset,
                ipod_directory[i].chksum,
                ipod_directory[i].vers,
-               ipod_directory[i].loadAddr);
+               ipod_directory[i].loadAddr,
+               ipod_directory[i].devOffset+sector_size+((ipod_directory[i].len+sector_size-1)&~(sector_size-1)));
     }
 #endif
 
@@ -448,14 +687,19 @@ int list_images(int nimages, struct ipod_directory_t* ipod_directory)
          switch(ipod_directory[i].ftype) {
              case FTYPE_OSOS:
                  if (ipod_directory[i].entryOffset==0) {
-                     printf("    Main firmware - %d bytes\n",ipod_directory[i].len);
+                     printf("    Main firmware - %d bytes\n",
+                       ipod_directory[i].len);
                  } else {
-                     printf("    Main firmware - %d bytes\n",ipod_directory[i].entryOffset);
-                     printf("    Third-party bootloader - %d bytes\n",ipod_directory[i].len-ipod_directory[i].entryOffset);
+                     printf("    Main firmware - %d bytes\n",
+                       ipod_directory[i].entryOffset);
+                     printf("    Third-party bootloader - %d bytes\n",
+                       ipod_directory[i].len-ipod_directory[i].entryOffset);
                  }
                  break;
              default:
-                     printf("    %s - %d bytes\n",ftypename[ipod_directory[i].ftype],ipod_directory[i].len);
+                     printf("    %s - %d bytes\n",
+                            ftypename[ipod_directory[i].ftype],
+                            ipod_directory[i].len);
          }
     }
     printf("\n");
@@ -463,12 +707,21 @@ int list_images(int nimages, struct ipod_directory_t* ipod_directory)
     return 0;
 }
 
+
+
+
 int main(int argc, char* argv[])
 {
     int i;
+    int infile, outfile;
     int ipod_version;
+    unsigned int inputsize;
     struct partinfo_t pinfo[4]; /* space for 4 partitions on 1 drive */
     int nimages;
+    off_t diroffset;
+    char* modelname;
+    int modelnum;
+    char* filename;
     struct ipod_directory_t ipod_directory[MAX_IMAGES];
     int action = SHOW_INFO;
     int sector_size;
@@ -498,6 +751,37 @@ int main(int argc, char* argv[])
         if ((strcmp(argv[i],"-l")==0) || (strcmp(argv[i],"--list")==0)) {
             action = LIST_IMAGES;
             i++;
+        } else if (strcmp(argv[i],"--remove-bootloader")==0) {
+            action = REMOVE_BOOTLOADER;
+            i++;
+        } else if ((strcmp(argv[i],"-ef")==0) || 
+                   (strcmp(argv[i],"--extract-firmware")==0)) {
+            action = EXTRACT_FIRMWARE;
+            i++;
+            if (i == argc) { print_usage(); return 1; }
+            filename=argv[i];
+            i++;
+        } else if ((strcmp(argv[i],"-rf")==0) || 
+                   (strcmp(argv[i],"--replace-firmware")==0)) {
+            action = REPLACE_FIRMWARE;
+            i++;
+            if (i == argc) { print_usage(); return 1; }
+            filename=argv[i];
+            i++;
+        } else if ((strcmp(argv[i],"-r")==0) || 
+                   (strcmp(argv[i],"--read-partition")==0)) {
+            action = READ_PARTITION;
+            i++;
+            if (i == argc) { print_usage(); return 1; }
+            filename=argv[i];
+            i++;
+        } else if ((strcmp(argv[i],"-w")==0) || 
+                   (strcmp(argv[i],"--write-partition")==0)) {
+            action = WRITE_PARTITION;
+            i++;
+            if (i == argc) { print_usage(); return 1; }
+            filename=argv[i];
+            i++;
         } else {
             print_usage(); return 1;
         }
@@ -531,7 +815,8 @@ int main(int argc, char* argv[])
         return 3;
     }
 
-    nimages=read_directory(dh, pinfo[0].start*sector_size, sector_size, ipod_directory);
+    nimages=read_directory(dh, pinfo[0].start*sector_size, sector_size, 
+                           ipod_directory, &diroffset);
     if (nimages <= 0) {
         fprintf(stderr,"[ERR]  Failed to read firmware directory\n");
         return 1;
@@ -540,30 +825,93 @@ int main(int argc, char* argv[])
     ipod_version=(ipod_directory[0].vers>>12) & 0x0f;
     printf("[INFO] Ipod model: ");
     switch (ipod_version) {
-        case 0x3: printf("3rd Generation\n"); break;
-        case 0x4: printf("1st Generation Mini\n"); break;
-        case 0x5: printf("4th Generation\n"); break;
-        case 0x6: printf("Photo/Color\n"); break;
-        case 0x7: printf("2nd Generation Mini\n"); break;
-        case 0xc: printf("1st Generation Nano\n"); break;
-        case 0xb: printf("Video (aka 5th Generation)\n"); break;
-        default: printf("UNKNOWN (Firmware version is %08x)\n",ipod_directory[0].vers);
+        case 0x3:
+            printf("3rd Generation\n");
+            modelnum = 7;
+            modelname = "ip3g";
+            break;
+        case 0x4:
+            printf("1st Generation Mini\n");
+            modelnum = 9;
+            modelname = "mini";
+            break;
+        case 0x5:
+            printf("4th Generation\n");
+            modelnum = 8;
+            modelname = "ip4g";
+            break;
+        case 0x6:
+            printf("Photo/Color\n");
+            modelnum = 3;
+            modelname = "ipco";
+            break;
+        case 0x7:
+            printf("2nd Generation Mini\n");
+            modelnum = 11;
+            modelname = "mn2g";
+            break;
+        case 0xc:
+            printf("1st Generation Nano\n");
+            modelnum = 4;
+            modelname = "nano";
+            break;
+        case 0xb:
+            printf("Video (aka 5th Generation)\n");
+            modelnum = 5;
+            modelname = "ipvd";
+            break;
+        default:
+            printf("[ERR] Unknown firmware version (0x%08x)\n",
+                   ipod_directory[0].vers);
+            return -1;
     }
 
     if (action==LIST_IMAGES) {
-        list_images(nimages,ipod_directory);
-#if 0
-    } else if (mode==READ) {
+        list_images(nimages,ipod_directory,sector_size);
+    } else if (action==REMOVE_BOOTLOADER) {
+        if (ipod_directory[0].entryOffset==0) {
+            fprintf(stderr,"[ERR]  No bootloader detected.\n");
+        } else {
+            if (remove_bootloader(dh, pinfo[0].start*sector_size, sector_size, 
+                                  ipod_directory)==0) {
+                fprintf(stderr,"[INFO] Bootloader removed.\n");
+            }
+        }
+    } else if (action==REPLACE_FIRMWARE) {
+        if (ipod_reopen_rw(&dh, devicename) < 0) {
+            return 5;
+        }
+
+        if (replace_firmware(dh, filename,pinfo[0].start*sector_size, 
+                             sector_size, nimages, ipod_directory, diroffset, 
+                             modelnum, modelname)==0) {
+            fprintf(stderr,"[INFO] Firmware replaced with %s.\n",filename);
+        } else {
+            fprintf(stderr,"[ERR]  --replace-firmware failed.\n");
+        }
+    } else if (action==EXTRACT_FIRMWARE) {
+        if (extract_firmware(dh, filename,pinfo[0].start*sector_size, 
+                             sector_size, ipod_directory, modelnum, modelname
+                            )==0) {
+            fprintf(stderr,"[INFO] Firmware extracted to %s.\n",filename);
+        } else {
+            fprintf(stderr,"[ERR]  --extract-firmware failed.\n");
+        }
+    } else if (action==READ_PARTITION) {
         outfile = open(filename,O_CREAT|O_WRONLY|O_BINARY,S_IREAD|S_IWRITE);
         if (outfile < 0) {
            perror(filename);
            return 4;
         }
 
-        res = disk_read(dh,outfile,pinfo[p].start,pinfo[p].size,sector_size);
-
+        if (read_partition(dh, outfile, pinfo[0].start, pinfo[0].size, 
+                           sector_size) < 0) {
+            fprintf(stderr,"[ERR]  --read-partition failed.\n");
+        } else {
+            fprintf(stderr,"[INFO] Partition extracted to %s.\n",filename);
+        }
         close(outfile);
-    } else if (mode==WRITE) {
+    } else if (action==WRITE_PARTITION) {
         if (ipod_reopen_rw(&dh, devicename) < 0) {
             return 5;
         }
@@ -577,16 +925,20 @@ int main(int argc, char* argv[])
         /* Check filesize is <= partition size */
         inputsize=filesize(infile);
         if (inputsize > 0) {
-            if (inputsize <= (pinfo[p].size*sector_size)) {
-                fprintf(stderr,"[INFO] Input file is %lu bytes\n",inputsize);
-                res = disk_write(dh,infile,pinfo[p].start,sector_size);
+            if (inputsize <= (pinfo[0].size*sector_size)) {
+                fprintf(stderr,"[INFO] Input file is %u bytes\n",inputsize);
+                if (write_partition(dh,infile,pinfo[0].start,
+                    sector_size) < 0) {
+                    fprintf(stderr,"[ERR]  --write-partition failed.\n");
+                } else {
+                    fprintf(stderr,"[INFO] %s restored to partition\n",filename);
+                }
             } else {
                 fprintf(stderr,"[ERR]  File is too large for firmware partition, aborting.\n");
             }
         }
 
         close(infile);
-#endif
     }
 
     ipod_close(dh);