#! /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

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)

def ownDevice(blockSize = 4194304):
    """
    Find the device path of the USB stick itself, when cloning in "own" mode
    
    @param blockSize the size of wanted blocks; defaults to 4194304 == 4M
    @return the path of the USB stick's device and its size in blocks unit
    """
    try:
        mountpoint = glob.glob("/usr/lib/live/mount/persistence/sd*1")[0]
    except:
        mountpoint = "/usr/lib/live/mount/persistence/sda1"
    device = re.sub(
        r"/usr/lib/live/mount/persistence/(sd.)1", r"/dev/\1", mountpoint)
    cmd=f"sfdisk -l {device}| grep /dev/sd.1| awk '{{print $4}}'"
    completed = run(cmd, shell=True, capture_output=True, encoding="utf-8")
    sectors=int(completed.stdout.strip())
    count = 1 + sectors*512//blockSize # convert 512 B sectors to 4 MiB blocks
    return device, count

class StickMaker(QObject):
    def __init__(self):
        QObject.__init__(self)
        self.startTime = time.time()
        self.pathSentence="PATH=/usr/sbin:/sbin:/usr/bin:/bin "

    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):
        """
        calls a command, with PATHS set via self.pathSentence,
        then call self.flush()
        @param cmd a shell command
        @returncode the child return code.
        """
        returncode = call(self.pathSentence+cmd, shell=True)
        self.flush()
        return returncode
        

    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', help=self.tr('the targetted USB stick'))
        
        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:
            self.file = self.device
            completed = run("/sbin/losetup --show --find "+ self.file,
                            shell=True, capture_output=True, encoding="utf-8")
            time.sleep(1)
            self.loopdevice = completed.stdout.strip()
            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)
        

        ################### 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()
            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)

        # in case there are partitions #4 and #3 in the table, delete them
        for num in (4, 3):
            completed = run(
                f"{self.pathSentence} sfdisk -l {self.device}| grep {self.device}{num}",
                shell=True, capture_output=True, encoding="utf-8"
            )
            if completed.stdout.strip():
                self.log(self.tr("Delete the partition {}{}, to create it again").format(self.device, num))
                run(f"{self.pathSentence} sfdisk --delete {self.device} {num}", shell=True)
        ############# eject and reload the USB stick to let    #######
        ############# the kernel understand its new structure  #######
        self.log(self.tr("Eject and reload the disk {}, to let the kernel understand its new structure").format(self.device))
        run(["sudo", "sg_start", "--eject", self.device])
        time.sleep(1)
        run(["sudo", "sg_start", "--load", self.device])
        
        ############# adds partition #3 for the VFAT data   ##########
        if self.simulation:
            self.log(self.tr("Create the partition {}p3 for VFAT data, using {}GiB").format(self.device, self.vfat_size))
        else:
            self.log(self.tr("Create the partition {}3 for VFAT data, using {}GiB").format(self.device, self.vfat_size))

        def do_createVfat():
            size = int(float(self.vfat_size)*1024)
            cmd = f"{self.pathSentence} echo ',{size}M,0c' | {self.pathSentence} sfdisk -N3 {self.device} --append"
            completed = run(
                cmd, shell=True, capture_output=True, encoding="utf-8")
            self.sync()
            return

        def test_createVfat():
            cmd = f"{self.pathSentence} sfdisk -l {self.device}"
            completed = run(
                cmd, shell=True, capture_output=True, encoding="utf-8")
            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_createVfat, do_createVfat, comment = "Create Vfat")
        
        ############# eject and reload the USB stick to let    #######
        ############# the kernel understand its new structure  #######
        self.log(self.tr("Eject and reload the disk {}, to let the kernel understand its new structure").format(self.device))
        run(["sudo", "sg_start", "--eject", self.device])
        time.sleep(1)
        run(["sudo", "sg_start", "--load", self.device])

        ############# adds partition #4 for the persistence ##########
        if self.simulation:
            self.log(self.tr("Create the partition {}p4 to support persistence").format(self.device))
        else:
            self.log(self.tr("Create the partition {}4 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():
            cmd = f"{self.pathSentence} echo ','|sfdisk -N4 {self.device} --append"
            completed = run(
                cmd, shell=True, capture_output=True, encoding="utf-8")
            self.log(completed.stdout)
            self.log(completed.stderr)
            self.sync()
            ### the command /usr/sbin/fdisk seems to pace down the kernel ###
            self.call_cmd(f'printf "p\nq\n" | /sbin/fdisk {self.device}')
            return

        def test_createExt4():
            cmd = f"{self.pathSentence} sfdisk -l {self.device}"
            completed = run(
                cmd, shell=True, capture_output=True, encoding="utf-8")
            lines = completed.stdout.split("\n")
            if self.simulation:
                found = [l for l in lines if l.startswith(self.device+"p4")]
            else:
                found = [l for l in lines if l.startswith(self.device+"4")]
            return bool(found)
            
        ensure (test_createExt4, do_createExt4, comment = "Create Ext4")

        ############# eject and reload the USB stick to let    #######
        ############# the kernel understand its new structure  #######
        self.log(self.tr("Eject and reload the disk {}, to let the kernel understand its new structure").format(self.device))
        run(["sudo", "sg_start", "--eject", self.device])
        time.sleep(1)
        run(["sudo", "sg_start", "--load", self.device])

        def do_formats():
            ################ format VFAT   ####################################
            if self.simulation:
                self.call_cmd(f"mkfs.vfat -n DATA {self.device}p3")
            else:
                self.call_cmd(f"mkfs.vfat -n DATA {self.device}3")
            ################ format (ext4) ####################################
            if self.simulation:
                self.call_cmd(f"mkfs.ext4 -L persistence -F {self.device}p4")
            else:
                self.call_cmd(f"mkfs.ext4 -L persistence -F {self.device}4")
            return
        
        def test_formats():
            if self.simulation:
                return self.call_cmd(f"fsck.vfat -a {self.device}p3") == 0 and \
                self.call_cmd(f"fsck.ext4 -p {self.device}p4") == 0
            return self.call_cmd(f"fsck.vfat -a {self.device}3") == 0 and \
                self.call_cmd(f"fsck.ext4 -p {self.device}4") == 0
        
        ensure(test_formats, do_formats, comment="Format vfat and ext4 partitions")
        ############# eject and reload the USB stick to let    #######
        ############# the kernel understand its new structure  #######
        self.log(self.tr("Eject and reload the disk {}, to let the kernel understand its new structure").format(self.device))
        run(["sudo", "sg_start", "--eject", self.device])
        time.sleep(1)
        run(["sudo", "sg_start", "--load", self.device])

        ################## mount the 4th partition ########################

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

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

        self.log(self.tr("Create the file persistence.conf on the fourth partition"))
        if self.simulation:
             with open(f"/tmp/{self.device_short}p4/persistence.conf", "w") as outfile:
                 outfile.write("/ union\n")
        else:
            with open(f"/tmp/{self.device_short}4/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 fourth partition with rw.tar.gz"))
            if self.simulation:
                self.call_cmd(f"zcat rw.tar.gz | (cd /tmp/{self.device_short}p4/; tar xf -)")
            else:
                self.call_cmd(f"zcat rw.tar.gz | (cd /tmp/{self.device_short}4/; tar xf -)")

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

        if self.simulation:
            self.call_cmd(f"umount /tmp/{self.device_short}p4; rmdir /tmp/{self.device_short}p4; partprobe {self.device}")
        else:
            self.call_cmd(f"umount /tmp/{self.device_short}4; rmdir /tmp/{self.device_short}4; partprobe {self.device}")
        self.sync(delay=0)

        self.call_cmd(f"sfdisk -l {self.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}
loop = {self.loopdevice}""")
        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()
