qwad/WiiPy/archive.py

455 lines
13 KiB
Python
Raw Normal View History

2012-08-21 19:28:14 +02:00
from common import *
from title import *
import zlib
class U8(WiiArchive):
"""This class can unpack and pack U8 archives, which are used all over the Wii. They are often used in Banners and contents in Downloadable Titles. Please remove all headers and compression first, kthx.
The f parameter is either the source folder to pack, or the source file to unpack."""
class U8Header(Struct):
__endian__ = Struct.BE
def __format__(self):
self.tag = Struct.string(4)
self.rootnode_offset = Struct.uint32
self.header_size = Struct.uint32
self.data_offset = Struct.uint32
self.zeroes = Struct.string(16)
class U8Node(Struct):
__endian__ = Struct.BE
def __format__(self):
self.type = Struct.uint16
self.name_offset = Struct.uint16
self.data_offset = Struct.uint32
self.size = Struct.uint32
def __init__(self):
self.files = []
def _dump(self):
header = self.U8Header()
rootnode = self.U8Node()
# constants
header.tag = "U\xAA8-"
header.rootnode_offset = 0x20
header.zeroes = "\x00" * 16
rootnode.type = 0x0100
nodes = []
strings = "\x00"
data = ''
for item, value in self.files:
node = self.U8Node()
recursion = item.count('/')
if(recursion < 0):
recursion = 0
name = item[item.rfind('/') + 1:]
node.name_offset = len(strings)
strings += name + '\x00'
if(value == None):
node.type = 0x0100
node.data_offset = recursion
node.size = len(nodes)
for one, two in self.files:
if(one[:len(item)] == item): # find nodes in the folder
node.size += 1
node.size += 1
else:
sz = len(value)
node.data_offset = len(data)
data += value + "\x00" * (align(sz, 32) - sz) # 32 seems to work best for fuzzyness? I'm still really not sure
node.size = sz
node.type = 0x0000
nodes.append(node)
header.header_size = ((len(nodes) + 1) * len(rootnode)) + len(strings)
header.data_offset = align(header.header_size + header.rootnode_offset, 64)
rootnode.size = len(nodes) + 1
for i in range(len(nodes)):
if(nodes[i].type == 0x0000):
nodes[i].data_offset += header.data_offset
fd = ''
fd += header.pack()
fd += rootnode.pack()
for nodeobj in nodes:
fd += nodeobj.pack()
fd += strings
fd += "\x00" * (header.data_offset - header.rootnode_offset - header.header_size)
fd += data
return fd
def _dumpDir(self, dir):
if(not os.path.isdir(dir)):
os.mkdir(dir)
old = os.getcwd()
os.chdir(dir)
for item, data in self.files:
if(data == None):
if(not os.path.isdir(item)):
os.mkdir(item)
else:
open(item, "wb").write(data)
os.chdir(old)
def _loadDir(self, dir):
try:
self._tmpPath += ''
except:
self._tmpPath = ''
old = os.getcwd()
os.chdir(dir)
entries = os.listdir(".")
for entry in entries:
if(os.path.isdir(entry)):
self.files.append((self._tmpPath + entry, None))
self._tmpPath += entry + '/'
self._loadDir(entry)
elif(os.path.isfile(entry)):
data = open(entry, "rb").read()
self.files.append((self._tmpPath + entry, data))
os.chdir(old)
self._tmpPath = self._tmpPath[:self._tmpPath.find('/') + 1]
def _load(self, data):
offset = 0
header = self.U8Header()
header.unpack(data[offset:offset + len(header)])
offset += len(header)
assert header.tag == "U\xAA8-"
offset = header.rootnode_offset
rootnode = self.U8Node()
rootnode.unpack(data[offset:offset + len(rootnode)])
offset += len(rootnode)
nodes = []
for i in range(rootnode.size - 1):
node = self.U8Node()
node.unpack(data[offset:offset + len(node)])
offset += len(node)
nodes.append(node)
strings = data[offset:offset + header.data_offset - len(header) - (len(rootnode) * rootnode.size)]
offset += len(strings)
recursion = [rootnode.size]
recursiondir = []
counter = 0
for node in nodes:
counter += 1
name = strings[node.name_offset:].split('\0', 1)[0]
if(node.type == 0x0100): # folder
recursion.append(node.size)
recursiondir.append(name)
assert len(recursion) == node.data_offset + 2 # haxx
self.files.append(('/'.join(recursiondir), None))
elif(node.type == 0): # file
self.files.append(('/'.join(recursiondir) + '/' + name, data[node.data_offset:node.data_offset + node.size]))
offset += node.size
else: # unknown type -- wtf?
pass
sz = recursion.pop()
if(sz != counter + 1):
recursion.append(sz)
else:
recursiondir.pop()
def __str__(self):
ret = ''
for key, value in self.files:
name = key[key.rfind('/') + 1:]
recursion = key.count('/')
ret += ' ' * recursion
if(value == None):
ret += '[' + name + ']'
else:
ret += name
ret += '\n'
return ret
def __getitem__(self, key):
for item, val in self.files:
if(item == key):
if(val != None):
return val
else:
ret = []
for item2, val2 in self.files:
if(item2.find(item) == 0):
ret.append(item2[len(item) + 1:])
return ret[1:]
raise KeyError
def __setitem__(self, key, val):
for i in range(len(self.files)):
if(self.files[i][0] == key):
self.files[i] = (self.files[i][0], val)
return
self.files.append((key, val))
class WAD(WiiArchive):
def __init__(self, boot2 = False):
self.tmd = TMD()
self.tik = Ticket()
self.contents = []
self.boot2 = False
self.cert = ""
def _load(self, data):
if(self.boot2 != True):
headersize, wadtype, certsize, reserved, tiksize, tmdsize, datasize, footersize, padding = struct.unpack('>I4s6I32s', data[:64])
pos = 64
else:
headersize, data_offset, certsize, tiksize, tmdsize, padding = struct.unpack('>IIIII12s', data[:32])
pos = 32
rawcert = data[pos:pos + certsize]
pos += certsize
if(self.boot2 != True):
if(certsize % 64 != 0):
pos += 64 - (certsize % 64)
self.cert = rawcert
rawtik = data[pos:pos + tiksize]
pos += tiksize
if(self.boot2 != True):
if(tiksize % 64 != 0):
pos += 64 - (tiksize % 64)
self.tik = Ticket.load(rawtik)
rawtmd = data[pos:pos + tmdsize]
pos += tmdsize
if(self.boot2 == True):
pos = data_offset
else:
pos += 64 - (tmdsize % 64)
self.tmd = TMD.load(rawtmd)
titlekey = self.tik.getTitleKey()
contents = self.tmd.getContents()
for i in range(0, len(contents)):
tmpsize = contents[i].size
if(tmpsize % 16 != 0):
tmpsize += 16 - (tmpsize % 16)
encdata = data[pos:pos + tmpsize]
pos += tmpsize
decdata = Crypto().decryptContent(titlekey, contents[i].index, encdata)
self.contents.append(decdata)
if(tmpsize % 64 != 0):
pos += 64 - (tmpsize % 64)
def _loadDir(self, dir):
origdir = os.getcwd()
os.chdir(dir)
self.tmd = TMD.loadFile("tmd")
self.tik = Ticket.loadFile("tik")
self.cert = open("cert", "rb").read()
contents = self.tmd.getContents()
for i in range(len(contents)):
self.contents.append(open("%08x.app" % i, "rb").read())
os.chdir(origdir)
def _dumpDir(self, dir):
origdir = os.getcwd()
os.chdir(dir)
contents = self.tmd.getContents()
for i in range(len(contents)):
open("%08x.app" % i, "wb").write(self.contents[i])
self.tmd.dumpFile("tmd")
self.tik.dumpFile("tik")
open("cert", "wb").write(self.cert)
os.chdir(origdir)
def _dump(self, fakesign = True):
titlekey = self.tik.getTitleKey()
contents = self.tmd.getContents()
apppack = ""
for i, content in enumerate(contents):
if(fakesign):
content.hash = str(Crypto().createSHAHash(self.contents[content.index]))
content.size = len(self.contents[content.index])
encdata = Crypto().encryptContent(titlekey, content.index, self.contents[content.index])
apppack += encdata
if(len(encdata) % 64 != 0):
apppack += "\x00" * (64 - (len(encdata) % 64))
if(fakesign):
self.tmd.setContents(contents)
self.tmd.fakesign()
self.tik.fakesign()
rawtmd = self.tmd.dump()
rawcert = self.cert
rawtik = self.tik.dump()
sz = 0
for i in range(len(contents)):
sz += contents[i].size
if(sz % 64 != 0):
sz += 64 - (contents[i].size % 64)
if(self.boot2 != True):
pack = struct.pack('>I4s6I', 32, "Is\x00\x00", len(rawcert), 0, len(rawtik), len(rawtmd), sz, 0)
pack += "\x00" * 32
else:
pack = struct.pack('>IIIII12s', 32, align(len(rawcert) + len(rawtik) + len(rawtmd), 0x40), len(rawcert), len(rawtik), len(rawtmd), "\x00" * 12)
pack += rawcert
if(len(rawcert) % 64 != 0 and self.boot2 != True):
pack += "\x00" * (64 - (len(rawcert) % 64))
pack += rawtik
if(len(rawtik) % 64 != 0 and self.boot2 != True):
pack += "\x00" * (64 - (len(rawtik) % 64))
pack += rawtmd
if(len(rawtmd) % 64 != 0 and self.boot2 != True):
pack += "\x00" * (64 - (len(rawtmd) % 64))
if(self.boot2 == True):
pack += "\x00" * (align(len(rawcert) + len(rawtik) + len(rawtmd), 0x40) - (len(rawcert) + len(rawtik) + len(rawtmd)))
pack += apppack
return pack
def __getitem__(self, idx):
return self.contents[idx]
def __setitem__(self, idx, value):
self.contents[idx] = value
def __str__(self):
out = ""
out += "Wii WAD:\n"
out += str(self.tmd)
out += str(self.tik)
return out
class CCF():
class CCFHeader(Struct):
__endian__ = Struct.LE
def __format__(self):
self.magic = Struct.string(4)
self.zeroes12 = Struct.string(12)
self.rootOffset = Struct.uint32
self.filesCount = Struct.uint32
self.zeroes8 = Struct.string(8)
class CCFFile(Struct):
__endian__ = Struct.LE
def __format__(self):
self.fileName = Struct.string(20)
self.fileOffset = Struct.uint32
self.fileSize = Struct.uint32
self.fileSizeDecompressed = Struct.uint32
def __init__(self, fileName):
self.fileName = fileName
self.fd = open(fileName, 'r+b')
def compress(self, folder):
fileList = []
fileHdr = self.CCFHeader()
files = os.listdir(folder)
fileHdr.magic = "\x43\x43\x46\x00"
fileHdr.zeroes12 = '\x00' * 12
fileHdr.rootOffset = 0x20
fileHdr.zeroes8 = '\x00' * 8
currentOffset = len(fileHdr)
packedFiles = 0
previousFileEndOffset = 0
for file in files:
if os.path.isdir(folder + '/' + file) or file == '.DS_Store':
continue
else:
fileList.append(file)
fileHdr.filesCount = len(fileList)
self.fd.write(fileHdr.pack())
self.fd.write('\x00' * (fileHdr.filesCount * len(self.CCFFile())))
for fileNumber in range(len(fileList)):
fileEntry = self.CCFFile()
compressedBuffer = zlib.compress(open(folder + '/' + fileList[fileNumber]).read())
fileEntry.fileName = fileList[fileNumber]
fileEntry.fileSize = len(compressedBuffer)
fileEntry.fileSizeDecompressed = os.stat(folder + '/' + fileList[fileNumber]).st_size
fileEntry.fileOffset = align(self.fd.tell(), 32) / 32
print 'File {0} ({1}Kb) placed at offset 0x{2:X}'.format(fileEntry.fileName, fileEntry.fileSize / 1024, fileEntry.fileOffset * 32)
self.fd.seek(len(fileHdr) + fileNumber * len(self.CCFFile()))
self.fd.write(fileEntry.pack())
self.fd.seek(fileEntry.fileOffset * 32)
self.fd.write(compressedBuffer)
self.fd.close()
def decompress(self):
fileHdr = self.CCFHeader()
hdrData = self.fd.read(len(fileHdr))
fileHdr.unpack(hdrData)
print 'Found {0} file/s and root node at 0x{1:X}'.format(fileHdr.filesCount, fileHdr.rootOffset)
if fileHdr.magic != "\x43\x43\x46\x00":
raise ValueError("Wrong magic, 0x{0}".format(fileHdr.magic))
try:
os.mkdir(os.path.dirname(self.fileName) + '/' + self.fd.name.replace(".", "_") + "_out")
except:
pass
os.chdir(os.path.dirname(self.fileName) + '/' + self.fd.name.replace(".", "_") + "_out")
currentOffset = len(fileHdr)
for x in range(fileHdr.filesCount):
self.fd.seek(currentOffset)
fileEntry = self.CCFFile()
fileData = self.fd.read(len(fileEntry))
fileEntry.unpack(fileData)
fileEntry.fileOffset = fileEntry.fileOffset * 32
print 'File {0} at offset 0x{1:X}'.format(fileEntry.fileName, fileEntry.fileOffset)
print 'File size {0}Kb ({1}Kb decompressed)'.format(fileEntry.fileSize / 1024, fileEntry.fileSizeDecompressed / 1024)
output = open(fileEntry.fileName.rstrip('\0'), 'w+b')
self.fd.seek(fileEntry.fileOffset)
if fileEntry.fileSize == fileEntry.fileSizeDecompressed:
print 'The file is stored uncompressed'
output.write(self.fd.read(fileEntry.fileSize))
else:
print 'The file is stored compressed..decompressing'
decompressedBuffer = zlib.decompress(self.fd.read(fileEntry.fileSize))
output.write(decompressedBuffer)
output.close()
currentOffset += len(fileEntry)
if(__name__ == '__main__'):
wad = WAD.loadFile("testing.wad")
print wad
wad.dumpDir("outdir")
wad.dumpFile("interesting.wad", fakesign = False) #keyword arguements work as expected when calling _dump(). awesome.
wad2 = WAD.loadDir("outdir")
print wad2
wad3 = WAD.loadFile("interesting.wad")
print wad3
wad3.dumpDir("outdir2")