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

# wbfs2fat.py - WBFS to FAT32 conversion
# (c) 2012 by PsyBlade - all rights reserved
# http://gbatemp.net/t291320-wbfs2fat-py
version = "0.9.0beta2"
date = "2013-01-03"

if __name__ == "__main__":
    import wbfs2fat, sys
    wbfs2fat.main()
    sys.exit(0)

import ctypes, io, logging, re, struct, time
import argparse, logging.handlers, subprocess, sys, textwrap
import cfg, wbfs, fat

logger = logging.getLogger(__name__)

def findparts():
    if sys.platform in ["linux2","cygwin"]:
        partfile = open("/proc/partitions","rt")
        partfile.readline()
        partfile.readline()
        parts=[("/dev/"+line.split()[3], "/dev/"+line.split()[3]) for line in partfile]
        partfile.close()
    elif sys.platform == "win32":
        import wmi
        w32 = wmi.WMI ()
        parts=[("\\\\.\\"+disk.DeviceID,disk.DeviceID) for disk in w32.Win32_LogicalDisk()]
    return parts


class WbfsError(Exception):
    pass


class Status(dict):
    #normal
    info_wbfs   = (0,"INFO: {} WBFS with {:>3} Games")
    #auto fixable
    err_xlink   = (3,'ERROR: WBFS is corrupted: Blocks used multiple times (try "wwt help verify" to fix or use ignore-xlink to ignore)')
    #manually fixable
    err_nofree  = (4,"ERROR: Aditional free space needed: {} ({}x{} {})")
    #fatal
    err_nomagic = (5,"ERROR: Not a WBFS partition (missing magic)")
    err_unkvers = (5,"ERROR: Unknown WBFS Version: {}")
    err_segv    = (5,'ERROR: WBFS is corrupted: Blocks used are outside filesystem boundry (try deleting them to fix)')
    err_size    = (5,'ERROR: WBFS is corrupted: Filesystem is larger than device')
    #manually fixable
    err_io      = (6,"IOERROR {}: {}")
    #status
    stat_noscan = (7,"INFO: not scanned")
    info = set(); info.add(info_wbfs)

    def less(self, *key):
        return max(self.keys()) < key

    def maxmaj(self):
        return max(self.keys())[0]

    def convertable(self):
        return self.less(4,0)

    def writeable(self):
        return self.less(6,0)

    def scanned(self):
        return self.less(7,0)

    def add(self, key, *var):
        if not self.setdefault(key,var) == var: raise KeyError
        return key[1].format(*var)

    def rm(self, key):
        self.pop(key, None)

    def onestr(self, key):
        return key[1].format(*self[key])

    def tostr(self):
        return "\n".join(key[1].format(*var) for key, var in sorted(self.items()))


class Device:
    def __init__(self, device, name=None):
        self.name = name or device
        self.dev_name = device
        self.status = Status([(Status.stat_noscan,"")])
        self.opened = ""

    def open_ro(self):
        self.dev = io.open(self.dev_name, "rb")
        self.opened = "rb"

    def open_rw(self):
        self.dev = io.open(self.dev_name, "r+b")
        self.opened = "r+b"

    def close(self):
        if self.opened:
            self.dev.close()
            self.opened = ""

    def scan(self):
        if cfg.options.convert: cfg.report(1,"checking for WBFS")
        try:
            self.open_ro()
            self.wbfs = wbfs.Wbfs(self.dev.peek(12), self.status)
            if cfg.options.convert: cfg.report (1,"analysing WBFS")
            self.wbfs.raw = self.dev.read(1<<self.wbfs.block_size_p)
            self.wbfs.analyse(self.dev)
        except WbfsError as e:
            logger.debug("Exception caught",exc_info=1)
        except IOError as e:
            logger.debug("Exception caught",exc_info=1)
            if e[0]==22 and sys.platform == "win32":
                if self.name == self.dev_name:
                    self.close()
                    self.dev_name = "\\\\.\\"+self.dev_name
                    self.scan()
                else:
                    self.dev_name = self.name
                    self.status.add(Status.err_io, *e[:2])
            else:
                self.status.add(Status.err_io, *e[:2])
        finally:
            self.close()
            self.status.rm(Status.stat_noscan)

    def backup(self):
        cfg.report(0,"creating Backup of overlap")
        cfg.backup_file = cfg.options.backup_dir+"/wbfs."+re.sub(":","",self.dev_name.strip().split('\\')[-1].split('/')[-1]+"."+str(int(time.time())))+".back"
#        cfg.report(2," "+cfg.backup_file)
        with io.BufferedWriter(io.FileIO(cfg.backup_file,"wb")) as back:
            self.dev.seek(0)
            for i in range(self.wbfs.fat_blocks):
#                cfg.report (3," "+str(i))
                back.write(self.dev.read(self.wbfs.block_size))

    def restore(self):
        cfg.report(0,"restoring Backup")
        with io.BufferedReader(io.FileIO(cfg.backup_file,"rb")) as back:
            self.dev.seek(0)
            block = back.read(1<<20)
            i = 0
            while len(block)>0:
#                cfg.report (3," "+str(i))
                i+=1
                self.dev.write(block)
                block = back.read(1<<20)
        cfg.report(0,"restore done")

    def manual_restore(self):
        cfg.confirm("ready to restore {} to {}".format(cfg.options.restore, cfg.dev_name))
        cfg.backup_file = cfg.options.restore
        self.open_rw()
        self.restore()
        self.close()

    def create_mbr(self):
        cfg.report(1,"creating MBR")
        self.mbr=ctypes.create_string_buffer(self.wbfs.block_size)
        #0x0000 440 Bootstrap code area
        #0x01b8 4 Disk signature (optional)
        #0x01bc 2 usually nulls; 0x0000
        #0x01be 1 p1 bootable
        #0x01bf 3 p1 CHS start
        sf2B.pack_into(self.mbr,0x01bf,0xff)
        sf2B.pack_into(self.mbr,0x01c0,0xff)
        sf2B.pack_into(self.mbr,0x01c1,0xff)
        #0x01c2 1 p1 type
        sf2B.pack_into(self.mbr,0x01c2,0x0b)
        #0x01c3 3 p1 CHS end
        sf2B.pack_into(self.mbr,0x01c3,0xff)
        sf2B.pack_into(self.mbr,0x01c4,0xff)
        sf2B.pack_into(self.mbr,0x01c5,0xff)
        #0x01c6 4 p1 LBA start
        sf2I.pack_into(self.mbr,0x01c6,self.wbfs.block_sects)
        #0x01ca 4 p1 LBA length
        sf2I.pack_into(self.mbr,0x01ca,self.wbfs.sect_cnt - self.wbfs.block_sects)
        #0x01ce 16 p2
        #0x01de 16 p3
        #0x01ee 16 p4
        #0x01FE 2 MBR signature 0x55 0xAA
        sf2H.pack_into(self.mbr,0x01FE,0x0AA55)

    def write_mbr(self):
        cfg.report(1,"writing MBR to Disk")
        self.dev.seek(0)
        self.dev.write(self.mbr)

    def init_fat(self):
        self.fat=fat.Fat(self.wbfs)

    def write_head(self):
        cfg.report(1,"writing WBFS File Headers")
        for game in self.wbfs.games:
            game.write_head(self.dev)

    def write_relocate(self):
        cfg.report(1,"relocating overlaping Blocks")
        dev = self.dev
        block2adr = self.wbfs.block2adr
        for (old,  _), new in sorted(self.wbfs.remap.items()):
#            cfg.report(3,"\t"+str(old)+"->"+str(new))
            dev.seek(old<<self.wbfs.vblock_size_p)
            data = dev.read(self.wbfs.vblock_size)
            dev.seek(new<<self.wbfs.vblock_size_p)
            dev.write(data)

    def write_fat(self):
        cfg.report(1,"writing FAT to Disk")
        self.dev.seek(cfg.options.create_mbr * self.wbfs.block_size)
        for i in range(self.wbfs.fat_blocks):
#            cfg.report (3," "+str(i))
            self.dev.write(self.fat.raw[i*self.wbfs.block_size:(i+1)*self.wbfs.block_size])

    def fsck(self):
        if sys.platform == "linux2":
            cfg.report(1,"checking FAT for Errors")
            fsckcmd=['fsck.vfat', '-a', self.dev_name]
            if cfg.options.verbosity>=4: fsckcmd.append("-l")
            if cfg.options.verbosity>=3: fsckcmd.append("-v")
            fsck = subprocess.Popen(fsckcmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            res = fsck.communicate()
            ret = fsck.returncode
            if ret or res[1]:
#                if res[0]: cfg.report(2,"\n"+res[0].decode("ascii"))
#                if res[1]: cfg.report(3,"\n"+res[1].decode("ascii"))
                fsck = subprocess.Popen(fsckcmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
                res = fsck.communicate()
                ret = fsck.returncode
#            if res[0]: cfg.report(4,"\n"+res[0].decode("ascii"))
#            if res[1]: cfg.report(0,"\n"+res[1].decode("ascii"))
            if ret or res[1]:
                raise WbfsError(3,0,'fsck error {}: {}'.format(ret, res[1]))

    def ntfs(self):
        if sys.platform == "win32":
            cfg.report(1,"converting to NTFS")
            if subprocess.call(['convert', self.dev_name, "/x", "/nosecurity", "/fs:ntfs"]): sys.exit()

    def convert(self):
        if not max(self.status.viewkeys()-cfg.ignore) < (1,0): raise WbfsError()

        self.init_fat()
        cfg.report(1,"creating FAT filesystem")
        self.fat.mkfs()
        cfg.report(1,"creating Files")
        for game in self.wbfs.games:
            game.mkfiles(self.fat.wbfsdir)
        if cfg.options.create_mbr: self.create_mbr()

        self.open_rw()
        if cfg.options.backup: self.backup()
        cfg.confirm("ready for conversion")
        cfg.report(1,"WARNING: start of dangerous part")
        try:
            self.write_head()
            if self.wbfs.remap: self.write_relocate()
            self.write_fat()
            if cfg.options.create_mbr: self.write_mbr()
        except Exception as e:
            logger.error("Exception caught",exc_info=1)
            if cfg.options.auto_restore:
#                self.status = 3
                self.restore()
#                self.status = 8
            else:
#                self.status = 5
                raise
        finally:
            self.close()

        if max(self.status.viewkeys()-cfg.ignore) < (1,0) and cfg.options.fsck:
            try:
                self.fsck()
            except Exception as e:
                logger.error("Exception caught",exc_info=1)
                if cfg.options.auto_restore:
#                    self.status = 3
                    self.open_rw()
                    self.restore()
                    self.close()
#                    self.status = 8
                else:
#                    self.status = 5
                    raise

        if max(self.status.viewkeys()-cfg.ignore) < (1,0):
            cfg.report(1,"Conversion completed")
        else: print self.status.tostr()


    def checkoptions(self, fix=0):
        abort = 0
        msg = []
        if cfg.options.convert:
#            msg = [self.status.onestr(Status.info_wbfs)]
            if not cfg.options.backup:  msg.append("WARNING: skipping backup! Only do this if you desperatly short on disk space.")
            elif not cfg.options.backup_dir:
                if fix: abort |= 1; cfg.options.backup_dir="."
                else:   abort |= 2; msg.append("ERROR: Backup Directory is required for backup! Please select one.")
            if cfg.options.auto_restore:
                if not cfg.options.backup:
                    if fix: abort |= 1; cfg.options.auto_restore = 0
                    else:   abort |= 2; msg.append("ERROR: Automatic restore needs backup creation enabled")
            else:   msg.append("WARNING: Drive might be left in unusable state on error! Use manual restore if that happens.")
            if cfg.options.fsck:
                if not sys.platform == "linux2":
                    if fix: abort |= 1; cfg.options.fsck = 0
                    else:   abort |= 2; msg.append("ERROR: fsck curently supported on Linux")
                if cfg.options.create_mbr:
                    if fix: abort |= 1; cfg.options.fsck = 0
                    else:   abort |= 2; msg.append("ERROR: fsck with create MBR curently not supported")
            if Status.err_xlink in self.status:
                if not Status.err_xlink in cfg.ignore and fix: abort |= 1; cfg.ignore.add(Status.err_xlink)
                if Status.err_xlink in cfg.ignore:  msg.append("WARNING: Manually overriding crosslink abort! Multiply used blocks will be duplicated. Read the documentation on what this means.")
            if cfg.options.create_mbr:   msg.append("WARNING: Only create an MBR if you are absolutely sure that you use an unpartitioned drive! "+
                    "While possible thats really uncommon. Using a single partition is much more common. "+
                    "You might need to reinsert it after conversion.")

            for key in self.status.viewkeys()-cfg.ignore-Status.info:
                abort |= 2
                msg.append(self.status.onestr(key))

        elif cfg.options.restore != None:
            msg.append("WARNING: Restoring a Backup File after writing to the filesystem is dangerous! Try to avoid if possible.")
            if not cfg.options.restore:   abort |= 4; msg.append("ERROR: Backup File is required for restore! Please select one.")
        else: raise Exception
        return abort, msg


def main():
    log = logging.getLogger()
    log.setLevel(logging.DEBUG)
    fh = logging.handlers.RotatingFileHandler('wbfs2fat.log', backupCount=5, delay=1)
    try:
        fh.doRollover()
    except OSError:
        pass
    #fh.setLevel(logging.DEBUG)
    fh.setFormatter(logging.Formatter('%(levelname)s %(name)s %(asctime)s: %(message)s'))
    log.addHandler(fh)
    ch = logging.StreamHandler()
    ch.setLevel(logging.ERROR)
    ch.setFormatter(logging.Formatter('%(levelname)s %(name)s: %(message)s'))
    log.addHandler(ch)

#    usage = "usage: %prog [options] device"
    parser = argparse.ArgumentParser()#usage,version = "%prog "+version)
    parser.add_argument("-v", "--verbose", action="count", dest="verbosity", default=1)
    parser.add_argument("-q", "--quiet", action="store_const", const=0, dest="verbosity")
    parser.add_argument("-Y", "--yes", action="store_false", dest="confirm", default=1, help="asume yes on all questions (DANGEROUS)")
    parser.add_argument("dev_names", metavar="DEV", nargs="*", help="devices to operate on")

    mode_group = parser.add_argument_group("operation mode: (default: gui)")
    mode = mode_group.add_mutually_exclusive_group()
    mode.add_argument("--gui", action="store_true", help="start gui and scan devices")
    mode.add_argument("--scan", action="store_true", help="scan devices for WBFS, leave empty for auto")
    mode.add_argument("--convert", action="store_true", help="convert a device to FAT32")
    mode.add_argument("--restore", metavar="BACKUP", help="manually restore a backup to a device")

    convert = parser.add_argument_group("options for convert & gui")
    convert.add_argument("--backup-dir", default=".", metavar="DIR", help="directory to store backup in")
    convert.add_argument("-B", "--no-backup", action="store_false", dest="backup", default=1, help="do not create a backup of the changed data")
    convert.add_argument("-R", "--no-auto-restore", action="store_false", dest="auto_restore", default=1, help="do not automatically restore the backup in case of error")
    convert.add_argument("-F", "--no-fsck", action="store_false", dest="fsck", default=1, help="do not run fsck on linux")
    convert.add_argument("-s", "--scrub", action="store_true", dest="scrub", default=0, help="enable scrubbing, saves some space but is slower")
    convert.add_argument("-X", "--ignore-xlink", nargs=0, action=cfg.Ignore, const=Status.err_xlink, help="do not abort if blocks are used multiple times, duplicate them instead")
    convert.add_argument("-M", "--create-mbr", action="store_true", default=0,  help="create a mbr to fix unpartitioned drives (HIGHLY EXPERIMENTAL AND DANGEROUS)")

    cfg.options = parser.parse_args()
    if len(cfg.options.dev_names) == 1:
        cfg.dev_name = cfg.options.dev_names[0]
    if not sys.platform == "linux2":
        cfg.options.fsck = 0

    if cfg.options.scan:
        import cli
        ch.setLevel(logging.ERROR-10*cfg.options.verbosity)
        plist = zip(cfg.options.dev_names,cfg.options.dev_names) or findparts()
        wrapper_init = textwrap.TextWrapper(subsequent_indent="\t"*4)
        wrapper_subs = textwrap.TextWrapper(initial_indent="\t"*3, subsequent_indent="\t"*4)
        for part in plist:
            dev = Device(*part)
            dev.scan()
            out = "{:<23} {}".format(dev.name, dev.status.tostr()).splitlines()
            cfg.report(0,wrapper_init.fill(out[0]))
            for line in out[1:]:
                cfg.report(0,wrapper_subs.fill(line))
        print

    elif cfg.options.convert or cfg.options.restore:
        if cfg.dev_name:
            import cli
            ch.setLevel(logging.ERROR-10*cfg.options.verbosity)
            wrapper = textwrap.TextWrapper(subsequent_indent="\t")
            try:
                dev = Device(cfg.dev_name)
                dev.scan()
                abort, check = dev.checkoptions(not cfg.options.confirm)
                for line in check:
                    cfg.report(0, wrapper.fill(line))
                if abort > 1:
                    print
                    sys.exit(1)
                if cfg.options.restore:
                    dev.manual_restore()
                else:
                    dev.convert()

            except Exception:
                logger.error("Exception caught",exc_info=1)
                sys.exit(1)

        else:
            cfg.report(0,"--convert and --restore need exactly one device to operate on")
            sys.exit(2)

    else:
        import gui.gui
        gui.gui.startgui()

    sys.exit(0)
