import os.path
import struct
from binascii import b2a_hex
from hexdump import hexdump, asasc, tohex
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
"""

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

        self.readtad()

    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))
        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''
            if dokodecode and not args.nokod:
                plain = dat[:plainbytes]
                dat = kodecode(i+1, dat[plainbytes:])
            if args.ascdump:
                print("%5d: %08x-%08x: (%02x:%08x) %s %s" % (i, ofs, ofs+ln, flags, chk, tohex(plain), asasc(dat)))
            else:
                print("%5d: %08x-%08x: (%02x:%08x) %s %s" % (i, ofs, ofs+ln, flags, chk, tohex(plain), 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 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)
    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()
        decode_kod(args, data)


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

    if db.stru:
        print("stru")
        db.stru.dump(args, dokodecode=True)
    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 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('--shift', '-s', type=str)
    ko.add_argument('--ascdump', '-a', action='store_true')
    ko.add_argument('--nokod', '-n', action='store_true')
    ko.add_argument('filename', type=str, nargs='?')
    ko.set_defaults(handler=kod_hexdump)

    cro = subparsers.add_parser('crodump', help='CROdumper')
    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('dbdir', type=str)
    cro.set_defaults(handler=cro_dump)

    args = parser.parse_args()

    if args.handler:
        args.handler(args)


if __name__=='__main__':
    main()