import os.path
import io
import struct
from binascii import b2a_hex
from hexdump import hexdump, asasc, tohex, unhex
from koddecoder import kodecode
"""
python3 crodump.py crodump chechnya_proverki_ul_2012
python3 crodump.py kodump -s 6   -o 0x4cc9 -e 0x5d95 chechnya_proverki_ul_2012/CroStru.dat
"""

def enumunreferenced(ranges, filesize):
    o = 0
    for start, end, desc in sorted(ranges):
        if start > o:
            yield o, start-o
        o = end
    if o<filesize:
        yield o, filesize-o

class Datafile:
    def __init__(self, dat, tad):
        self.dat = dat
        self.tad = tad

        self.readtad()

        self.dat.seek(0, io.SEEK_END)
        self.datsize = self.dat.tell()

    def readtad(self):
        self.tad.seek(0)
        hdrdata = self.tad.read(2*4)
        self.tadhdr = struct.unpack("<2L", hdrdata)
        indexdata = self.tad.read()
        self.tadidx = [ struct.unpack_from("<3L", indexdata, 12*_) for _ in range(len(indexdata)//12) ]

    def readdata(self, ofs, size):
        self.dat.seek(ofs)
        return self.dat.read(size)

    def dump(self, args, dokodecode=False, plainbytes=0):
        print("tadhdr: %08x %08x" % tuple(self.tadhdr))
        ranges = []
        for i, (ofs, ln, chk) in enumerate(self.tadidx):
            if ln==0xFFFFFFFF:
                print("%5d: %08x %08x %08x" % (i, ofs, ln, chk))
                continue
            flags = ln>>24

            ln &= 0xFFFFFFF
            dat = self.readdata(ofs, ln)
            plain = b''
            decrypted = ' '
            if dokodecode and not args.nokod:
                pb = plainbytes if flags else 8
                plain = dat[:pb]
                dat = kodecode(i+1, dat[pb:])
                decrypted = '*' if flags else '+'
            if args.ascdump:
                print("%5d: %08x-%08x: (%02x:%08x) %s %s%s" % (i, ofs, ofs+ln, flags, chk, tohex(plain), decrypted, asasc(dat)))
            else:
                print("%5d: %08x-%08x: (%02x:%08x) %s %s%s" % (i, ofs, ofs+ln, flags, chk, tohex(plain), decrypted, tohex(dat)))
            ranges.append((ofs, ofs+ln, "item #%d" % i))

        if args.verbose:
            # output parts not referenced in the .tad file.
            for o, l in enumunreferenced(ranges, self.datsize):
                dat = self.readdata(o, l)
                if args.ascdump:
                    print("%08x-%08x: %s" % (o, o+l, asasc(dat)))
                else:
                    print("%08x-%08x: %s" % (o, o+l, tohex(dat)))


class Database:
    def __init__(self, dbdir):
        self.dbdir = dbdir

        self.stru = self.getfile("Stru")
        self.index = self.getfile("Index")
        self.bank  = self.getfile("Bank")
        self.sys = self.getfile("Sys")

    def getfile(self, name):
        try:
            return Datafile(open(self.getname(name, "dat"), "rb"), open(self.getname(name, "tad"), "rb"))
        except IOError:
            return

    def getname(self, name, ext):
        return os.path.join(self.dbdir, "Cro%s.%s" % (name, ext))

def incdata(data, s):
    """
    add 's' to each byte.
    This is useful for finding the correct shift from an incorrectly shifted chunk.
    """
    return b"".join(struct.pack("<B", (_+s)&0xFF) for _ in data)

def decode_kod(args, data):
    """
    various methods of hexdumping KOD decoded data.
    """
    if args.nokod:
        # plain hexdump, no KOD decode
        hexdump(args.offset, data)
    elif args.shift:
        # explicitly specified shift.
        args.shift = int(args.shift, 0)
        enc = kodecode(args.shift, data)
        hexdump(args.offset, enc)
    elif args.increment:
        # explicitly specified shift.
        for s in range(256):
            enc = incdata(data, s)
            if args.ascdump:
                print("%02x: %s" % (s, asasc(enc)))
            else:
                print("%02x: %s" % (s, tohex(enc)))



    else:
        # output with all possible 'shift' values.
        for s in range(256):
            enc = kodecode(s, data)
            if args.ascdump:
                print("%02x: %s" % (s, asasc(enc)))
            else:
                print("%02x: %s" % (s, tohex(enc)))



def kod_hexdump(args):
    """
    KOD decode a section of a data file
    """
    args.offset = int(args.offset, 0)
    if args.length:
        args.length = int(args.length, 0)
    elif args.endofs:
        args.endofs = int(args.endofs, 0)
        args.length = args.endofs - args.offset

    if args.filename:
        with open(args.filename, "rb") as fh:
            fh.seek(args.offset)
            data = fh.read(args.length)
            decode_kod(args, data)
    else:
        # no filename -> read from stdin.
        import sys
        data = sys.stdin.buffer.read()
        if args.unhex:
            data = unhex(data)
        decode_kod(args, data)


def cro_dump(args):
    db = Database(args.dbdir)

    if db.stru:
        print("stru")
        db.stru.dump(args, dokodecode=True)
    if args.struonly:
        return
    if db.index:
        print("index")
        db.index.dump(args)
    if db.bank:
        print("bank")
        db.bank.dump(args)
    if db.sys:
        print("sys")
        db.sys.dump(args, dokodecode=True, plainbytes=8)


def main():
    import argparse
    parser = argparse.ArgumentParser(description='CRO hexdumper')
    subparsers = parser.add_subparsers()
    parser.set_defaults(handler=None)

    ko = subparsers.add_parser('kodump', help='KOD/hex dumper')
    ko.add_argument('--offset', '-o', type=str, default="0")
    ko.add_argument('--length', '-l', type=str)
    ko.add_argument('--endofs', '-e', type=str)
    ko.add_argument('--unhex', '-x', action='store_true', help="assume the input contains hex data")
    ko.add_argument('--shift', '-s', type=str, help="KOD decode with the specified shift")
    ko.add_argument('--increment', '-i', action='store_true', help="assume data is already KOD decoded, but with wrong shift -> dump alternatives.")
    ko.add_argument('--ascdump', '-a', action='store_true', help="CP1251 asc dump of the data")
    ko.add_argument('--nokod', '-n', action='store_true', help="don't KOD decode")
    ko.add_argument('filename', type=str, nargs='?', help="dump either stdin, or the specified file")
    ko.set_defaults(handler=kod_hexdump)

    cro = subparsers.add_parser('crodump', help='CROdumper')
    cro.add_argument('--verbose', '-v', action='store_true')
    cro.add_argument('--kodecode', '-k', action='store_true')
    cro.add_argument('--ascdump', '-a', action='store_true')
    cro.add_argument('--nokod', '-n', action='store_true')
    cro.add_argument('--struonly', action='store_true')
    cro.add_argument('dbdir', type=str)
    cro.set_defaults(handler=cro_dump)

    args = parser.parse_args()

    if args.handler:
        args.handler(args)


if __name__=='__main__':
    main()