gdishrink

#!/usr/bin/python
# -*- coding: utf-8 -*-

import sys, os, shutil
try:
    from cStringIO import StringIO
except ImportError:
    from StringIO import StringIO

def gdishrink(gdi):
    l = loadgdi(gdi)
    update_gdi(gdi)
    if l[0][0] == 3:
        gdishrink_three(l)
    else:
        gdishrink_cdda(l)

def gdishrink_three(l):
    print '3-Tracks GDI are unsupported for now...'
    pass

def gdishrink_cdda(l):
    nbt = int(l[0][0])
    modes = [2048 if l[i][3] == '2048' else 2352 for i in 1,3,nbt]
    fn = outname(l[3][4])

    with CdImage(l[3][4], modes[1]) as f:
        length = rFNES(f)
        f.seek(0,0)
        with open(fn, 'wb') as g:
            _copy_buffered(f, g, length = length)
    for j in l[1][4],l[nbt][4]:
        with CdImage(j, modes[0]) as f, open(outname(j),'wb') as g:
            _copy_buffered(f,g)



def loadgdi(gdi):
    with open(gdi) as g:
        l = [i.split() for i in g.readlines() if i.split()]
    if not int(l[3][1]) == 45000:
        raise AssertionError('Invalid gdi file: track03 LBA should be 45000')
    return l

def bakfile(fn):
    if not os.path.isfile(fn+'.bak'):
        shutil.copy(fn, fn+'.bak')
    if not os.path.isfile(fn+'.bak'):
        raise IOError('Could not backup {} properly!'.format(fn))

def backupGDI(gdi, l=None):
    bakfile(gdi)
    if l == None:
        l = loadgdi(gdi)
    nbt = int(l[0][0])
    bakfile(l[1][4])
    t3name = l[3][4]
    bakfile(t3name)
    if nbt > 3:
        txname = l[nbt][4]
        bakfile(txname)

def update_gdi(gdi):
    l=loadgdi(gdi)
    nbt = int(l[0][0])
    backupGDI(gdi, l=l)
    tracks = 1,3 if nbt == 3 else 1,3,nbt

    for t in tracks:
        l[t][3] = '2048'
        l[t][4] = outname(l[t][4])

    with open(gdi, 'wb') as f:
        f.write('\n'.join([i for i in [' '.join(j) for j in l]]))

def outname(i):
    o = i.replace('.bin','').replace('.BIN','')
    if not o[-4:] == '.iso':
        o += '.iso'
    return o


def rFNEB(f, b=2048, offset=None): # Reverse find non-empty block
    if offset == None:
        f.seek(0,2)
        offset = f.tell()
    empty = '\x00'*b
    loc = False
    while offset > 0:
        f.seek(offset-b if b<offset else 0)
        if f.read(b) == empty:
            offset -= b
        else:
            loc = f.tell()
            break
    return loc

def rFNES(f, offset=None):  # Adaptively Reverse find non-empty 2048 kB sector
    # Blocks are 20MB, 10MB, 1MB, 100kB, 10kB, 2kB.
    # The idea is to reduce seeking.
    b = [1024*int(i) for i in [2e4, 1e4, 1e3, 1e2, 1e1,2]]
    loc = offset
    for i in b:
        loc=rFNEB(f, b=i, offset=loc)
    return loc



# Class CdImage was taken from the gditools project and is gplv3 licensed
class CdImage(file):
    """
    Class that allows opening a 2352 or 2048 bytes/sector data cd track
    as a 2048 bytes/sector one.
    """
    def __init__(self, filename, mode = 'auto', *args, **kwargs):

        if mode == 'auto':
            if filename[-4:] == '.iso': mode = 2048
            elif filename[-4:] == '.bin': mode = 2352

        elif not mode in [2048, 2352]:
            raise ValueError('Argument mode should be either 2048 or 2352')
        self.__mode = mode

        if (len(args) > 0) and (args[0] not in ['r','rb']):
            raise NotImplementedError('Only read mode is implemented.')

        file.__init__(self, filename, 'rb')

        file.seek(self,0,2)
        if self.__mode == 2352:
            self.length = file.tell(self) * 2048/2352
        else:
            self.length = file.tell(self)
        file.seek(self,0,0)

        self.seek(0)

    def realOffset(self,a):
        return a/2048*2352 + a%2048 + 16

    def seek(self, a, b = 0):
        if self.__mode == 2048:
            file.seek(self, a, b)

        elif self.__mode == 2352:
            if b == 0:
                self.binpointer = a
            if b == 1:
                self.binpointer += a
            if b == 2:
                self.binpointer = self.length - a

            realpointer = self.realOffset(self.binpointer)
            file.seek(self, realpointer, 0)

    def read(self, length = None):
        if self.__mode == 2048:
            return file.read(self, length)

        elif self.__mode == 2352:
            if length == None:
                length = self.length - self.binpointer

            # Amount of bytes left until beginning of next sector
            tmp = 2048 - self.binpointer % 2048
            FutureOffset = self.binpointer + length
            realLength = self.realOffset(FutureOffset) - \
                            self.realOffset(self.binpointer)
            # This will (hopefully) accelerates readings on HDDs at the
            # cost of more memory use.
            buff = StringIO(file.read(self, realLength))
            # The first read can be < 2048 bytes
            data = buff.read(tmp)
            length -= tmp
            buff.seek(304, 1)
            # The middle reads are all 2048 so we optimize here!
            for i in xrange(length / 2048):
                data += buff.read(2048)
                buff.seek(304, 1)
            # The last read can be < 2048 bytes
            data += buff.read(length % 2048)
            # Seek back to where we should be
            self.seek(FutureOffset)
            return data

    def tell(self):
        if self.__mode == 2048:
            return file.tell(self)

        elif self.__mode == 2352:
            return self.binpointer

# Function _copy_buffered was taken from the gditools project and is gplv3 licensed
def _copy_buffered(f1, f2, length = False, bufsize = 1*1024*1024, closeOut = True):
    """
    Copy istream f1 into ostream f2 in bufsize chunks
    """
    if not length:  # By default it read all the file
        tmp = f1.tell()
        f1.seek(0,2)
        length = f1.tell()
        f1.seek(tmp,0)
    f2.seek(0,0)

    for i in xrange(length/bufsize):
        f2.write(f1.read(bufsize))
    f2.write(f1.read(length % bufsize))

    if closeOut:
        f2.close()



def main(argv):
    if len(argv) == 2:
        gdiname = argv[1]
        owd = os.getcwd()
        os.chdir(os.path.dirname(gdiname))
        gdishrink(os.path.basename(gdiname))
        os.chdir(owd)
    else:
        print "Usage:\n\tgdishrink relative/or/absolute/path/to/disc.gdi"


if __name__ == '__main__':
    main(sys.argv)