#This is for Python 2
from Crypto.Cipher import AES #pycrypto. If on windows, http://www.voidspace.org.uk/python/modules.shtml#pycrypto
import os, sys, struct
import hashlib

def roundUp(v, align):
    return ((v + align - 1) / align) * align

dsidevcommonkey = 'a1604a6a7123b529ae8bec32c816fcaa'.decode('hex')

if len(sys.argv) == 1:
    print 'Usage: %s poop.tad' % (os.path.basename(__file__))
    raise SystemExit()

with open(sys.argv[1], 'rb') as f:
    if f.read(6) != '\x00\x00\x00\x20\x49\x73':
        print 'Likely not a tad file'
        raise SystemExit()
    f.read(2)
    certsize = struct.unpack('>I', f.read(4))[0]
    crlsize = struct.unpack('>I', f.read(4))[0]
    ticketsize = struct.unpack('>I', f.read(4))[0]
    tmdsize = struct.unpack('>I', f.read(4))[0]
    contentsize = struct.unpack('>I', f.read(4))[0]
    metasize = struct.unpack('>I', f.read(4))[0]

    certOffset = roundUp(0x20, 64)
    crlOffset = roundUp(certOffset + certsize, 64)
    ticketOffset = roundUp(crlOffset + crlsize, 64)
    tmdOffset = roundUp(ticketOffset  + ticketsize, 64)
    contentOffset= roundUp(tmdOffset + tmdsize, 64)
    metaOffset = roundUp(contentOffset + contentsize, 64)
    fileSize = roundUp(metaOffset + metasize, 64)

    f.seek(ticketOffset+0x1BF) #encrypted title key
    enckey = f.read(16)
    f.seek(ticketOffset+0x1DC) #title id
    iv = f.read(8) + '\x00'*8
    titlekey = AES.new(dsidevcommonkey, AES.MODE_CBC, iv).decrypt(enckey)

    f.seek(tmdOffset+0x1DE)
    numcontents = struct.unpack('>H', f.read(2))[0]

    contents = []
    for i in xrange(numcontents): #parsing the tmd
        content = {}
        f.seek(tmdOffset+0x1E4+36*i)
        f.read(4) #content id
        content['idx'] = f.read(2) #content index
        f.read(2) #content type
        f.read(4) #first 4 bytes of content size. Not needed.
        content['size'] = struct.unpack('>I', f.read(4))[0]
        content['hash'] = f.read(20)
        contents.append(content)

    offset = 0
    for content in contents:
        f.seek(contentOffset+offset)
        encdata = f.read(content['size'])
        iv = '\x00'*14 + content['idx']
        decdata = AES.new(titlekey, AES.MODE_CBC, iv).decrypt(encdata)
        hash = hashlib.sha1(decdata).digest()
        if hash != content['hash']:
            print 'Hash mismatch! Decryption error!'
            raise SystemExit()
        with open('%s_dec_%d.srl' % (os.path.splitext(os.path.basename(sys.argv[1]))[0], struct.unpack('>H', content['idx'])[0]), 'wb') as outf:
            outf.write(decdata)
        offset += roundUp(content['size'], 64)