#!/usr/bin/python3

"""
Clone a Freeduc USB stick to the main hard disk: the target must be /dev/sda
"""

import glob, re, sys, time, tempfile
from subprocess import Popen, call, PIPE, run
from ensure import ensure
from ownsense import ownDevice

def error(msg):
    print(msg)
    print("Aborting.")
    sys.exit(1)
    return

def cloneISO():
    """
    Clone process when the usb stick bears an iso9660 file system
    """
    ## guess the size of the ISO image
    cmd = f"""sudo LANG=C /usr/sbin/sfdisk -l {sourceDevice} | grep '{sourceDevice}1'"""
    reply, _ = Popen(cmd, shell=True, stdout=PIPE, stderr=PIPE).communicate()
    m = re.match(r".* ([.\d]+)G.*", reply.decode())
    size=float(m.group(1)) # size in GB


    # erase the boot data
    call("dd if=/dev/zero of=/dev/sda bs=512 count=10", shell=True)

    ## remove swap space if any
    call("swapoff --all", shell=True)

    ## compose the command to clone to the hard disk
    cmd = f"/usr/bin/python3 /usr/lib/python3/dist-packages/live_clone/makeLiveStick.py --vfat-size 4 --no-persistence-seed --iso-filesize {size} --source own /dev/sda"
    # launch the command makeLiveStick.py
    call(cmd, shell=True)

    ## change the 5th partition to swap type
    cmd = "fdisk /dev/sda"
    p = Popen(cmd, shell = True, stdin = PIPE, stdout = PIPE)
    inputs = """\
    t
    5
    82
    w
    """
    p.stdin.write(inputs.encode("utf-8"))
    message,err=p.communicate()

    ## create the swap
    call("mkswap /dev/sda5", shell = True)

    ## write a mention in /etc/fstab
    call("mount /dev/sda6 /mnt", shell = True)
    call("mkdir -p /mnt/rw/etc", shell = True)
    call("mkdir -p /mnt/work", shell = True)
    call("cp /etc/fstab /mnt/rw/etc/", shell = True)
    call("echo '/dev/sda5 none swap sw 0 0' >>  /mnt/rw/etc/fstab",
         shell = True)
    call("umount /dev/sda6", shell = True)
    return

def cloneHDD(own_device):
    """
    Clone process when the usb stick bears a vfat file system
    @parameter own_device the device of the USB stick, e.g. '/dev/sdb'
    """
    cp = run(["/usr/sbin/fdisk", "-l", own_device],
             capture_output=True, encoding="utf-8")
    expr = own_device+r"1[^0-9]*(\d+)\s+(\d+)\s+(\d+).*"
    start, end, sectors = re.search(expr, cp.stdout, re.MULTILINE).groups()
    ## copying the <end> first sectors to the hard disk
    cmd = ["dd", f"if={own_device}", "of=/dev/sda", "bs=512",
           "count="+end, "status=progress"]
    call(cmd)
    call(["sync"])
    time.sleep(2)
    cmd = ["blockdev", "--rereadpt", "/dev/sda"]
    call(cmd)
    time.sleep(2)

    # modify the file /syslinux/live.cfg to erase occurences of
    # persistence-media=...
    cmd = ["mount", "/dev/sda1", "/mnt"]
    call(cmd)
    cmd = ["sed", "-i", "s/persistence-media=removable-usb //", "/mnt/syslinux/live.cfg"]
    call(cmd)
    cmd = ["umount", "/mnt"]
    call(cmd)
    
    ## commands for sfdisk:
    #### dump partitions, and forget everything after /dev/sda1
    cp = run(["sfdisk", "--dump", "/dev/sda"],
             capture_output=True, encoding="utf-8")
    structure = re.match(r".*/dev/sda1[^\n]*$",
                         cp.stdout, re.MULTILINE|re.DOTALL
                         ).group(0)+"\n"
    #### New Primary partition, number 2,  [default strat], 4 Gbytes, Swap
    #### New Primary partition, number 3, [default start], [default end], Linux
    structure += "/dev/sda2 : size=4G, type=S\n"
    structure += "/dev/sda3 : size=+, type=L\n"
    print("New structure =", structure)
    with tempfile.TemporaryFile("w", encoding="utf-8") as temp:
        temp.write(structure)
        temp.seek(0)
        cp = run(["sfdisk", "/dev/sda"], stdin=temp,
                 capture_output=True, encoding="utf-8")
        print(cp.stdout)
        print("errors:", cp.stderr)
    time.sleep(2)
    cmd = ["blockdev", "--rereadpt", own_device]
    call(cmd)
    time.sleep(2)

    call(["mkfs.ext4", "-L", "persistence", "/dev/sda3"])
    time.sleep(2)
    call(["mount", "/dev/sda3", "/mnt"])
    with open("/mnt/persistence.conf","w") as outfile:
        outfile.write("/ union\n")
    call(["umount", "/dev/sda3"])
    return

def go():
    """
    main program
    """
    ## check that we are running from a live Freeduc system
    partition1, _ = ownDevice()

    ## get the source device, maybe /dev/sdb
    sourceDevice=partition1[:-1]
    if sourceDevice == "/dev/sda":
        error(f"ERROR: Running {sys.argv[0]} from the hard disk is a nonsense.")

    start = time.time()
    cp = run("mount|grep medium",
             shell = True, capture_output=True, encoding="utf-8")
    if 'iso9660' in cp.stdout:
        cloneISO()
    elif 'vfat' in cp.stdout:
        cloneHDD(sourceDevice)
    else:
        error("ERROR: The source device's first partition must be "
              "either iso9660 or vfat")
    duration = int(time.time() - start)
    print(f"Cloning to hard disk finished, after {duration//60} minutes and {duration % 60:02d} seconds.")
    return

if __name__ == "__main__":
    go()
    
