#! /usr/bin/python3

import os, sys, argparse, glob, time, re
from subprocess import Popen, PIPE, call, run
from PyQt5.QtCore import QObject, QTranslator, QLocale
from PyQt5.QtWidgets import QApplication
from ensure import ensure
from ownsense import ownDevice

def disksize(device_short, disktype = "disk"):
    """
    Report the size of a disk, in bytes.
    
    @param device_short a name like "sdc" for device path /dev/sdc
    @param disktype someting like "disk" or "loop"; defaults to "disk"

    @return the size of the block device, in byte format, as an int.
    """
    progs = []
    progs.append("lsblk -b") # list block devices, size in byte format
    progs.append("grep "+device_short)
    if disktype == "disk":
        progs.append("grep disk")
    elif disktype == "loop":
        progs.append("grep loop")
        progs.append("grep -v part")
    progs.append("awk '{print $4}'")
    cmd = " | ".join(progs)
    completed = run(cmd, shell=True, capture_output=True, encoding="utf-8")
    size = completed.stdout.strip()
    return int(size)


class StickMaker(QObject):
    """
    A class to create a live stick.
    Parameters of the constructor
    @param verbose verbosity of messages, 0 by default
    """
    def __init__(self, verbose=0):
        QObject.__init__(self)
        self.startTime = time.time()
        self.pathSentence="PATH=/usr/sbin:/sbin:/usr/bin:/bin "
        self.verbose = verbose

    def hms(self):
        """
        returns the time elapsed so far, in format HH:MM:SS
        """
        t=int(time.time()-self.startTime)
        th=t//3600
        tmin=(t-3600*th)//60
        tsec=(t-3600*th-60*tmin)
        return f"{th:02d}:{tmin:02d}:{tsec:02d}"

    def flush(self, stderr=True, stdout=True):
        """
        flush standard outputs
        @param stderr don't flush sys.stderr if it's False (True by default)
        @param stdout don't flush sys.stdout if it's False (True by default)
        """
        if stderr: sys.stderr.flush()
        if stdout: sys.stdout.flush()
        return

    def call_cmd(self, cmd, min_verbosity = 2):
        """
        calls a command, with PATHS set via self.pathSentence,
        then call self.flush()
        @param cmd a shell command's template, with expressions surrounded
           by {} brackets, to be formatted with **self.__dict__
        @param min_verbosity a threshod to log a message, defaults to 2
        @returncode the child return code.
        """
        cmd = self.pathSentence + cmd.format(**self.__dict__)
        returncode = call(cmd, shell=True)
        self.flush()
        return returncode
        

    def run_cmd(self, cmd, min_verbosity = 2):
        """
        runs a command, with PATHS set via self.pathSentence,
        then calls self.flush()
        @param cmd a shell command's template, with expressions surrounded
           by {} brackets, to be formatted with **self.__dict__
        @param min_verbosity a threshod to log a message, defaults to 2
        @return a subprocess.CompletedProcess instance
        """
        cmd = self.pathSentence + cmd.format(**self.__dict__)
        if self.verbose >= min_verbosity:
            self.log(self.tr("Aiming to run the command {cmd}").format(cmd=cmd))
        completed = run(cmd, shell=True, capture_output=True, encoding="utf-8")
        self.flush()
        return completed

    def eject_reload(self, min_verbosity = 1):
        """
        ejects and reloads self.device
        @param min_verbosity a threshod to log a message, defaults to 1
        """
        if self.verbose >= min_verbosity:
            self.log(self.tr("Eject and reload the disk {}, to let the kernel understand its new structure").format(self.device))
        self.run_cmd("sg_start --eject {device}")
        time.sleep(1)
        self.run_cmd("sg_start --load {device}")
        return

    def log(self, s, timestamp=True):
        """
        print some information for stdout
        @param s a string to log
        @param timestamp True (by default) implies that there will be
        a timestamp
        """
        if timestamp:
            print(self.tr("[At {0}]    {1}").format(self.hms(),s))
        else:
            print(s)
        sys.stdout.flush()
        return

    def umountAll(self):
        """
        ensure all partitions of self.device are unmounted
        """
        completed = run(f"mount | grep {self.device} | awk '{{print $1}}'",
            shell=True, capture_output=True, encoding="utf-8"
        )
        mounts = completed.stdout.split()
        if mounts:
            run(["/usr/bin/umount"] + mounts)
        return
         
    def sync(self, delay = 2):
        """
        ensure that the target disk's structure is up to date with kernel's
        settings, and has no partitions mounted

        @param delay a time to sleep after sync and after partprobe; 
               defaults to 2 seconds
        """
        run(["/usr/bin/sync"])
        time.sleep(delay)

        self.umountAll()

        if self.simulation:
            run(["/sbin/losetup", "--detach", self.loopdevice])
            run(["/sbin/losetup", "--partscan", self.loopdevice, self.file])
        else:
            run(["/usr/sbin/blockdev","--rereadpt", self.device])
        time.sleep(delay) # dbus may mount partitions during a few seconds?
        return
        

    def stopUdisks(self):
        """
        Stop the Udidisks2 service if it is running; keep its previous state
        to launch it again later
        """
        completed = run("systemctl status udisks2",
                        shell=True, capture_output=True, encoding="utf-8")
        self.udisks_state = "active (running)" in completed.stdout
        if self.udisks_state:
            run("systemctl stop udisks2",
                shell=True, capture_output=True, encoding="utf-8")
        return
        
    def startUdisks(self):
        """
        Start the Udidisks2 service if it was running prior to self.go's call
        """
        if self.udisks_state:
            run("systemctl start udisks2",
                shell=True, capture_output=True, encoding="utf-8")
        return
        
    def go(self):
        """
        The main method
        """
        if os.geteuid() != 0:
            self.log(
                self.tr("You must be root to launch {}").format(sys.argv[0]),
                timestamp=False
            )
            sys.exit(1)


        parser = argparse.ArgumentParser(description=self.tr('clone a Debian-Live distribution'))
        parser.add_argument('-f','--vfat-size', help=self.tr("size of the VFAT partition, in GB"))
        parser.add_argument('-i','--iso-filesize', help=self.tr("size of the ISO image, in GB (required only when cloning from a block device)"))
        parser.add_argument('-n','--no-persistence-seed', help=self.tr("do'nt use the file rw.tar.gz to seed the persistence"), action='store_true')
        parser.add_argument('-s','--source', help=self.tr('source of the distribution (required); the keyword OWN means that we are cloning from a Debian-Live disk'), metavar="source", required=True)
        parser.add_argument('-m', '--simulation', help=self.tr('simulate the cloning onto a loop-mounted temporary file'), required=False, action='store_true')
        parser.add_argument('device', default="", nargs="?", help=self.tr('the targetted USB stick, or the name of the file to create in case of a simulation'))
        
        self.args = parser.parse_args()
        
        self.vfat_size = self.args.vfat_size
        if not self.vfat_size:
            self.vfat_size=1
        self.iso_filesize = self.args.iso_filesize
        self.no_persistence_seed = self.args.no_persistence_seed
        self.source = self.args.source
        self.simulation = self.args.simulation
        self.device = self.args.device
        if self.simulation:
            count = 14*256 # 14G divided by 4M
            self.file = self.args.device
            self.log(f"Padding 14G with zeros in file {self.file }")
            # create self.file initially filled with zeros
            completed = run(
                f"dd if=/dev/zero of={self.file} status=progress bs=4M oflag=dsync count={count}",
                shell=True)
            self.log(f"Create a loop device with {self.file}")
            completed = run(f"/sbin/losetup --show --find {self.file}",
                            shell=True, capture_output=True, encoding="utf-8")
            time.sleep(1)
            self.loopdevice = completed.stdout.strip()
            self.log(f"The loop device is: {self.loopdevice}")
            self.device = self.loopdevice
        self.device_short = self.device.replace("/dev/","")
            

        ################### unmount targetted partitions if any #############
        self.sync(delay = 0)
        ################### check the sizes #################################
        if self.simulation:
            size = disksize(self.device_short, "loop")
        else:
            size = disksize(self.device_short, "disk")
        # convert size to giga-bytes
        size = size / (1024 ** 3)
        try:
            iso_size = int(float(self.iso_filesize))
        except:
            iso_size = os.stat(self.source).st_size//1024**3
        min_persistence = 1 # add at least 1 GB for the persistence

        if iso_size + min_persistence > size:
            self.log(self.tr("The disk is too small ({} GiB) to contain the ISO image ({} GiB) plus at least 1 GB for the persistence").format(size, iso_size))
            assert (size > iso_size + min_persistence) # to stop the process

        if iso_size + float(self.vfat_size) + min_persistence > size:
            new_vfat_size = size - iso_size - min_persistence
            self.log(self.tr("The size of the VFAT partition has been shrinked down to {} GB").format(new_vfat_size))
            self.vfat_size = str(new_vfat_size)
        

        # in case there are partitions with numbers > 2, delete them
        completed = self.run_cmd("sfdisk -l {device}| grep {device}")
        for num in (6, 5, 4, 3):
            for l in completed.stdout.split("\n"):
                if f"{self.device}{num}" in l:
                    if self.verbose > 0:
                        self.log(self.tr("Delete the partition {}{}, to create it again").format(self.device, num))
                    self.run_cmd("sfdisk --delete {device} " + str(num))
        self.eject_reload()

        ################### copy the ISO-hybrid image #######################

        self.log(self.tr("Copying {0} to {1} ...").format(self.source, self.device))

        if self.source.upper()=="OWN":
            ### when copying from own, the count of 4M blocks is computed
            ### from the last sector used by ownDevice()
            sourceDevice, count = ownDevice()
            sourceDevice = sourceDevice.replace("1","") # strip the part. number
            self.call_cmd(f"dd if={sourceDevice} of={self.device} status=progress bs=4M oflag=dsync count={count}")
        else: # not "OWN"
            ### the boot sector will provide information for two partitions.
            self.call_cmd(f"dd if={self.source} of={self.device} status=progress bs=4M oflag=dsync")
        
        self.sync(delay = 3)

        self.eject_reload()

        ############# adds partition #3 as an extended partition #####
        if self.simulation:
            self.log(self.tr("Create the extended partition {}p3").format(self.device))
        else:
            self.log(self.tr("Create the extended partition {}3").format(self.device))

        def do_createExtendedPartition():
            self.umountAll()
            self.run_cmd("echo ',,Ex' | sfdisk --force -N3 {device}")
            self.sync()
            return

        def test_ExtendedPartition():
            cmd = f"{self.pathSentence} sfdisk -l {self.device}"
            completed = self.run_cmd("sfdisk -l {device}")
            lines = completed.stdout.split("\n")
            if self.simulation:
                found = [l for l in lines if l.startswith(self.device+"p3")]
            else:
                found = [l for l in lines if l.startswith(self.device+"3")]
            return bool(found)

        ensure (test_ExtendedPartition, do_createExtendedPartition,
                comment = self.tr("Create Extended Partition"))
        
        ############# adds partition #5 for the VFAT data   ##########
        if self.simulation:
            self.log(self.tr("Create the partition {}p5 for VFAT data, using {}GiB").format(self.device, self.vfat_size))
        else:
            self.log(self.tr("Create the partition {}5 for VFAT data, using {}GiB").format(self.device, self.vfat_size))

        def do_createVfat():
            self.umountAll()
            size = int(float(self.vfat_size)*1024)
            cmd = f"echo ',{size}M,0c' | sfdisk --force -N5 {self.device}"
            completed = self.run_cmd(cmd)
            self.sync()
            return

        def test_createVfat():
            cmd = f"{self.pathSentence} sfdisk -l {self.device}"
            completed = self.run_cmd("fdisk -l {device}")
            lines = completed.stdout.split("\n")
            if self.simulation:
                found = [l for l in lines if l.startswith(self.device+"p5")]
            else:
                found = [l for l in lines if l.startswith(self.device+"5")]
            return bool(found)
        
        ensure (test_createVfat, do_createVfat,
                comment = self.tr("Create Vfat"))
        
        self.eject_reload()

        ############# adds partition #6 for the persistence ##########
        if self.simulation:
            self.log(self.tr("Create the partition {}p6 to support persistence").format(self.device))
        else:
            self.log(self.tr("Create the partition {}6 to support persistence").format(self.device))
        self.log(self.tr("This partition is adjusted to use all the remaining space on {}").format(self.device))

        def do_createExt4():
            self.umountAll()
            completed = self.run_cmd("echo ','|sfdisk --force -N6 {device}")
            self.log(completed.stdout)
            self.log(completed.stderr)
            self.sync()
            ### the command /usr/sbin/fdisk seems to pace down the kernel ###
            self.call_cmd('printf "p\nq\n" | /sbin/fdisk {device}')
            return

        def test_createExt4():
            cmd = f"{self.pathSentence} sfdisk -l {self.device}"
            completed = self.run_cmd("sfdisk -l {device}")
            lines = completed.stdout.split("\n")
            if self.simulation:
                found = [l for l in lines if l.startswith(self.device+"p6")]
            else:
                found = [l for l in lines if l.startswith(self.device+"6")]
            return bool(found)
            
        ensure (test_createExt4, do_createExt4,
                comment = self.tr("Create Ext4"))

        self.eject_reload()

        def do_formats():
            ################ format VFAT   ####################################
            if self.simulation:
                self.call_cmd("mkfs.vfat -n DATA {device}p5")
            else:
                self.call_cmd("mkfs.vfat -n DATA {device}5")
            ################ format (ext4) ####################################
            if self.simulation:
                self.call_cmd("mkfs.ext4 -L persistence -F {device}p6")
            else:
                self.call_cmd("mkfs.ext4 -L persistence -F {device}6")
            return
        
        def test_formats():
            if self.simulation:
                return self.call_cmd("fsck.vfat -a {device}p5") == 0 and \
                self.call_cmd("fsck.ext4 -p {device}p6") == 0
            return self.call_cmd("fsck.vfat -a {device}5") == 0 and \
                self.call_cmd("fsck.ext4 -p {device}6") == 0
        
        ensure(test_formats, do_formats,
               comment=self.tr("Format vfat and ext4 partitions"))

        self.eject_reload()

        ################## mount the 6th partition ########################

        if self.simulation:
            self.call_cmd("mkdir -p /tmp/{device_short}p6; mount {device}p6 /tmp/{device_short}p6")
        else:
            self.call_cmd("mkdir -p /tmp/{device_short}6; mount {device}6 /tmp/{device_short}6")

        ################# write the file persistence.conf ################

        self.log(self.tr("Create the file persistence.conf on the last partition"))
        if self.simulation:
             with open(f"/tmp/{self.device_short}p6/persistence.conf", "w") as outfile:
                 outfile.write("/ union\n")
        else:
            with open(f"/tmp/{self.device_short}6/persistence.conf", "w") as outfile:
                outfile.write("/ union\n")

        ################## eventually seed the persistence partition ######

        if self.no_persistence_seed==False and os.path.exists("rw.tar.gz"):
            self.log(self.tr("Pre-seed the last partition with rw.tar.gz"))
            if self.simulation:
                self.call_cmd("zcat rw.tar.gz | (cd /tmp/{device_short}p6/; tar xf -)")
            else:
                self.call_cmd("zcat rw.tar.gz | (cd /tmp/{device_short}6/; tar xf -)")

        #################### little cleanup ###############################

        if self.simulation:
            self.call_cmd("umount /tmp/{device_short}p6; rm -rf /tmp/{device_short}p6; partprobe {device}")
        else:
            self.call_cmd("umount /tmp/{device_short}6; rm -rf /tmp/{device_short}6; partprobe {device}")
        self.sync(delay=0)

        self.call_cmd("sfdisk -l {device}")

        self.log(
            "\n=======================================================\n",
            timestamp=False
        )
        self.log(
            self.tr("Ready. Total time: {}").format(self.hms()),
            timestamp=False
        )
        if self.simulation:
            print(f"""\
In simulation mode:
===================
file = {self.file}
detaching the loop device: {self.loopdevice}""")
            # detach the loop device
            cp = run(["/sbin/losetup", "-la"],
                     capture_output=True, encoding="utf-8")
            if self.loopdevice in cp.stdout:
                run(["/sbin/losetup", "-d", self.loopdevice])
            # withdraw root privilege
            theuser = os.environ.get("SUDO_USER", "root")
            run(["chown", f"{theuser}:{theuser}", self.file])
        return

if __name__=="__main__":
    app = QApplication(sys.argv)
    # i18n stuff
    locale = QLocale.system().name()
    translation="live_clone_{}.ts".format(locale)
    langPath=os.path.join(os.path.abspath(os.path.dirname(__file__)),"lang",translation)
    translator = QTranslator(app)
    translator.load(langPath)
    app.installTranslator(translator)
    #############
    os.environ["XDG_RUNTIME_DIR"] = "/tmp/runtime-root"
    maker = StickMaker()
    maker.go()