#!/usr/bin/env python3

# CPK file extractor (e.g. used as ressource packages in some games)

import sys, struct, os, zlib

class DCPK:
    def read_uint(self):
        return struct.unpack('<I',self.f.read(4))[0]

    def __init__(self, cpk, out_dir=''):
        self.f = open(cpk,'rb')
        self.out_dir = out_dir
        assert(self.f.read(4)==b'DCPK')
        self.block_count = self.read_uint()
        self.data_offset = self.read_uint()
        self.data_size = self.read_uint()
        self.B = []
        for _ in range(self.block_count):
            name_length = self.read_uint()
            name = self.f.read(name_length).decode().rstrip('\x00')
            offset = self.read_uint()
            size = self.read_uint()
            unknown_field = self.read_uint()  # ??
            uncompressed_size = self.read_uint()
            compressed = self.read_uint()  # 0 or 1
            self.B.append((name,offset,size,compressed))

    def print_list(self):
        print('%d files in package' % self.block_count)
        for (name,_,size,_) in self.B:
            print('%s (%dB)' % (name,size))

    def extract(self, name_size_filter):
        if not os.path.exists(self.out_dir):
            os.makedirs(self.out_dir)
        for (name,offset,size,compressed) in self.B:
            if name_size_filter(name,size):
                fout = open(os.path.join(self.out_dir,name.replace('/','_')),'wb')
                self.f.seek(self.data_offset+offset)
                data = self.f.read(size)
                if compressed:
                    data = zlib.decompress(data)
                fout.write(data)
                fout.close()

def main():
    if not 2<=len(sys.argv)<=3:
        sys.stderr.write('usage: %s in.cpk [out_dir]\n' % sys.argv[0])
        sys.exit(1)
    out_dir = sys.argv[2] if len(sys.argv)==3 else 'OUT'

    # filter to extract music in some games
    name_size_filter = lambda name,size: 'Music' in name and size > 600<<10
    P = DCPK(sys.argv[1],out_dir)
    #P.print_list()
    P.extract(name_size_filter)

main()