ungpkg - PS3 - PSP - 2GB fix

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <fcntl.h>
#include <string.h>
#include <assert.h>
#include <stdlib.h>
#include <unistd.h>

#include "tools.h"
#include "types.h"
#include "common.h"

#ifdef WIN32
#define MKDIR(x,y) mkdir(x)
#else
#define MKDIR(x,y) mkdir(x,y)
#endif

typedef void (*decrypt_callback) (u8*, u64, u64, u8*);

static FILE *pkg;
static u64 size;
static u64 offset;
static PKG_HEADER *header = NULL;
static u8 key[0x10];
static u8 pkg_type;
static u8 psp_key[0x10];

static decrypt_callback decrypt;


void decrypt_debug_block(u8 *data, u64 size, u64 register_offset, u8 *_key) {
        u64 i;
        u8 key[0x40];
        u8 bfr[0x1c];

        memset(key, 0, sizeof key);
        memcpy(key, header->digest, 8);
        memcpy(key + 0x08, header->digest, 8);
        memcpy(key + 0x10, header->digest + 0x08, 8);
        memcpy(key + 0x18, header->digest + 0x08, 8);

        wbe64(key + 0x38, be64(key + 0x38) + register_offset);

        sha1(key, sizeof key, bfr);

        for(i=0; i<size;i++) {
                if(i != 0 && (i % 16) == 0) {
                        wbe64(key + 0x38, be64(key + 0x38) + 1);
                        sha1(key, sizeof key, bfr);
                }
                data[i] ^= bfr[i & 0xf];
        }
}

void decrypt_retail_block(u8 *data, u64 size, u64 register_offset, u8 *_key) {
        u64 tmp;
        u8 iv[0x10];

        memcpy(iv, header->k_licensee, 0x10);

        tmp = be64(iv + 8) + register_offset /*1*/;
        wbe64(iv + 8, tmp);
        if (tmp == 0)
                wbe64(iv, be64(iv) + register_offset /*1*/);


        aes128ctr(_key, iv, data, size ,data);
}


static void unpack_pkg(void) {

        PKG_FILE_HEADER *file;
        u64 i;
        u64 n_files;
        u32 fname_len;
        u64 file_offset;
        u32 flags;
        u32 fname_off;
        u8 *data_block;
        u8 *file_block;
        char fname[255];

        n_files = be32( (u8*) &header->item_count);
        file = malloc(sizeof(PKG_FILE_HEADER));


        for ( i = 0; i <  n_files; i++ ) {

                fseek(pkg, offset + (i * 0x20), SEEK_SET);
                fread(file, 0x20, 1, pkg);

                if((pkg_type & 0x000000f) == 2)
                    decrypt((u8*)file , 0x20,  (offset + (i * 0x20) - offset) / 16, psp_key);
                else
                    decrypt((u8*)file , 0x20,  (offset + (i * 0x20) - offset) / 16, key);

                fname_off = be32((u8*)&file->filename_offset);
                fname_len = be32((u8*)&file->filename_size);
                file_offset = be64((u8*)&file->data_offset);
                size = be64((u8*)&file->data_size);
                flags = be32((u8*)&file->flags);

                fseek(pkg, fname_off + offset, SEEK_SET);
                file_block = malloc( fname_len );
                fread( file_block, fname_len, 1, pkg );

                if( (flags & 0xf0000000) == 0x90000000) {
                    decrypt(file_block, fname_len,  fname_off / 16, psp_key);
                } else {
                    decrypt(file_block, fname_len,  fname_off / 16, key);
                }

                memset(fname, 0, sizeof fname);
                strncpy(fname,(const char*) file_block, fname_len);

                fseek( pkg, file_offset + offset, SEEK_SET);
                data_block = malloc ( size );
                fread( data_block, size, 1, pkg);

                if( (flags & 0xf0000000) == 0x90000000) {
                     decrypt(data_block, size, file_offset / 16, psp_key);
                } else {
                     decrypt(data_block, size, file_offset / 16, key);
                }

                flags &= 0xff;
                if(flags == 4) {
                        MKDIR(fname, 0777);
                } else if (flags == 1 || flags == 3 || flags == 2 || flags == 6) {
                        memcpy_to_file( fname, data_block, size);
                } else {
                        fail("unknown flags: %08x", flags);
                }
        }
}


int main(int argc, char *argv[]) {
        char *dir;

        if (argc != 2 && argc != 3)
                fail("usage: ungpkg filename.pkg [target]");

        pkg = fopen( (const char*) argv[1], "rb");

        header = malloc(sizeof(PKG_HEADER));
        fread( header, sizeof(PKG_HEADER), 1, pkg);

        if (argc == 2) {
                dir = malloc(0x31);
                memset(dir, 0, 0x31);
                memset(dir, 0, 0x30);
                memcpy(dir, header->contentid, 0x30);
        } else {
                dir = argv[2];
        }

        offset = be64((u8*)&header->data_offset);
        size = be64((u8*)&header->data_size);
        u32 type = be32((u8*)&header->pkg_type);

        printf("%x\n", header->magic);
        return;

        if( type & 0x80000000 ) {

                if ((type & 0x0000000F) != 1 && (type & 0x0000000F) != 2)
                        fail("invalid pkg type: %x", type);

                pkg_type = type;

                if (key_get_simple("gpkg-key", key, 0x10) < 0)
                        fail("failed to load the package key.");

                if (key_get_simple("psp-gpkg-key", psp_key, 0x10) < 0)
                        fail("failed to load the psp package key.");

                decrypt = decrypt_retail_block;
        } else {
                decrypt = decrypt_debug_block;
        }

        MKDIR(dir, 0777);

        if (chdir(dir) != 0)
                fail("chdir(%s)", dir);

        unpack_pkg();
        fclose(pkg);

        return 0;
}