#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;
}