# coding:utf-8

#!/usr/bin/python3
#
#  Copyright (c) 2007-2009 Canonical Ltd.
#
#  Author: Oliver Grawert <ogra@ubuntu.com>
#
#  Modifications 2013 from papoteur <papoteur@mageialinux-online.org>
#  and Geiger David <david.david@mageialinux-online.org>
#
#  This program is free software; you can redistribute it and/or
#  modify it under the terms of the GNU General Public License as
#  published by the Free Software Foundation; either version 2 of the
#  License, or (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program; if not, write to the Free Software Foundation,
#  Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301, USA.

import gettext
import hashlib
import io
import logging

###########
# imports #
###########
# import locale
import os
import sys
import time
from subprocess import Popen, PIPE


class Dumper(object):
    def _do_write(self, source, target, b, backup_mode, uid, gid):
        # Writing step
        # Dump mode
        self.returncode = 0
        self.operation = True
        bs = 4096 * 128
        try:
            ifc = io.open(source, "rb", 1)
        except:
            message = _("Reading error.") + source
            logging.debug(message)
            self.return_state = False
            self.return_message = message
            self.finished.set()
            return
        else:
            try:
                ofc = io.open(target, "wb", 0)
            except:
                message = _(
                    "You don't have permission to write to the device {}"
                ).format(target)
                logging.error(message)
                self.return_state = False
                self.return_message = message
                self.finished.set()
                return
            else:
                self.operation = True
                steps = list(range(0, b + 1, int(b / 100)))
                steps.append(b)
                indice = 1
                written = 0
                ncuts = int(b / bs)
                while ncuts <= 100:
                    bs = int(bs / 2)
                    ncuts = int(b / bs)
                for i in range(0, int(ncuts) + 1):
                    try:
                        buf = ifc.read(bs)
                    except:
                        message = _("Reading error.")
                        logging.error(message)
                        self.return_state = False
                        self.return_message = message
                        self.finished.set()
                        return
                    try:
                        ofc.write(buf)
                    except:
                        message = _("Writing error.")
                        logging.debug(message)
                        self.return_state = False
                        self.return_message = message
                        self.finished.set()
                        return
                    written += len(buf)
                    if written > steps[indice]:
                        if indice % 1 == 0:
                            self._progress = indice
                        indice += 1
                        try:
                            os.fsync(ofc)
                        except:
                            message = _("Writing error.")
                            logging.debug(message)
                            self.return_state = False
                            self.return_message = message
                            self.finished.set()
                            return
                if backup_mode:
                    # Restore user as owner
                    os.chown(target, uid, gid)
                try:
                    ofc.close()
                except:
                    message = _("Writing error.")
                    logging.error(message)
                    self.return_message = message
                    self.return_state = False
                    return
            ifc.close()
            self._progress = 100
            self.return_state = True
            self.return_message = _("Writing terminated")
            logging.debug(self.return_message)
            self.finished.set()
            return

    def _check_write(self, target, source):
        logging.debug("Start checking")
        self.return_message = ""
        messages = []
        b = os.path.getsize(source)
        # Compute the sum from the written device
        steps = list(range(0, b + 1, int(b / 100)))
        steps.append(b)
        indice = 0
        checked = 0
        sha512func = hashlib.sha3_512()
        ncuts = int(b / 1024)
        with open(target, "rb") as f:
            for x in range(0, ncuts):
                block = f.read(1024)
                sha512func.update(block)
                if checked > steps[indice]:
                    self._progress = indice
                    indice += 1
                checked += 1024
            logging.debug("last read %d" % (b - ncuts * 1024))
            block = f.read(b - ncuts * 1024)
            sha512func.update(block)
            sha512sumcalc = sha512func.hexdigest().upper()
        self.sum_check = ""
        sum_type = "sha3"
        sum_file = f"{source}.{sum_type}"
        try:
            # Look for sum files in the same directory as source
            with open(sum_file, "r") as fs:
                # Read the sum in the file
                self.sum_check = (fs.readline()).split()[0]
            self.sum_file = True
        except:
            logging.info(_("Sum file {} not found\n").format(sum_file))
            self.sum_file = False
        self.return_state = True
        if self.sum_check == "":  # Can't get stored sum
            messages.append(_("SHA3 sum: {}").format(sha512sumcalc))
        #  compare the sums
        elif sha512sumcalc == self.sum_check:
            # ,  keep the bracket, this is the place for sum type
            messages.append(_("The {} sum check is OK").format(sum_type))
        else:
            messages.append(
                _("/!\\The computed and stored sums don't match")
            )
            self.return_state = False
        self._progress = 100
        self.return_message = "\n".join(messages)
        logging.info(self.return_message)
        self.finished.set()

    def udev_wait(self, operation):
        wait = Popen(["udevadm", "settle", "--timeout=15"], stderr=PIPE)
        wait.communicate()
        working = True
        while working:
            time.sleep(0.5)
            wait.poll()
            rc = wait.returncode
            if not (rc is None):
                wait = None
                working = False
        if rc != 0:
            self.return_state = False
            self.return_message = _("Timeout reached when {}".format(operation))
            return False
        return True

    def _do_persistence(self, target, label, key, fs_type):
        logging.debug("Start doing persistence partition")
        p = Popen(["fdisk", target], stdin=PIPE, stderr=PIPE)
        # commands:
        # n : new
        # p : primary partition
        # 3 : partition number (1-4)
        # default first sector
        # default last sector
        # w : write and quit
        outs, errs = p.communicate(input=b"n\np\n3\n\n\nw\n")
        working = True
        while working:
            time.sleep(0.5)
            p.poll()
            rc = p.returncode
            if rc is None:
                working = True
            else:
                process = None
                working = False
        if rc != 0:
            if rc == 1:
                # the error is  Re-reading the partition table failed.: Device or resource busy
                logging.error(
                    _("Error {} while doing persistent partition: {}").format(
                        rc, errs.decode("utf-8")
                    )
                )
                logging.error(_("Try reloading partition table"))
                p = Popen(["partprobe", target], stdin=PIPE, stderr=PIPE)
                outs, errs = p.communicate()
                working = True
                while working:
                    time.sleep(0.5)
                    p.poll()
                    rc = p.returncode
                    if not (rc is None):
                        process = None
                        working = False
                    else:
                        self.return_state = False
                        self.return_message = _(
                            "Unable to reload the partition table: {}"
                        ).format(errs.decode("utf-8"))
                        logging.error(self.return_message)
                        self.finished.set()
                        return
            else:
                self.return_state = False
                self.return_message = _(
                    "Error {} while doing persistent partition: {}"
                ).format(rc, errs.decode("utf-8"))
                logging.error(self.return_message)
                self.finished.set()
                return
        logging.debug("New partition created")
        self._progress = 25

        ## Wait for propagation of the info of new partition table
        #if not self.udev_wait("creating a new partition table"):
            #logging.error("Timeout")
            #self.return_state = False
            #self.finished.set()
            #return
        if key == "":
            # example mkfs.ext4 -L mgalive-persist /dev/sdf3
            self.return_message = _("Additional partition added. Formatting...")
            path = target + "3"
            # Format partition according to the fs_type specified
            fs_type = fs_type.lower()
            if fs_type == "fat32":
                cmd = ["mkdosfs", "-F", "32", "-n", label.upper()[:11], path]
            elif fs_type == "exfat":
                cmd = ["mkfs.exfat", "-n", label.upper()[:11], path]
            elif fs_type == "ntfs":
                cmd = ["mkntfs", "-f", "-L", label, path]
            elif fs_type == "ext4":
                cmd = [
                        "mkfs.ext4",
                        "-q",
                        "-F",
                        "-L",
                        label,
                        path,
                    ]
            logging.info(f"Executing {' '.join(cmd)}")
            process = Popen(cmd, stderr=PIPE)
            outs, errs = process.communicate()
            working = True
            while working:
                time.sleep(0.5)
                process.poll()
                rc = process.returncode
                if rc is None:
                    working = True
                else:
                    process = None
                    working = False
            if rc == 0:
                self.return_state = True
                self._progress = 100
                self.return_message = _("Persistent partition done")
            else:
                self.return_state = False
                self.return_message = _(
                    "Error {} while doing persistent partition: {}"
                ).format(rc, errs.decode("utf-8"))
                self.finished.set()
                return
            logging.info("Persistent partition done")
            if not self.udev_wait(_("formatting partition")):
                return
        else:
            # example cryptsetup luksFormat /dev/sdb3
            self.return_message = _("Persistent partition added. Encrypting...")
            base_target = os.path.basename(target) + "3"
            process = Popen(
                ["cryptsetup", "luksFormat", "-q", target + "3", "-d", "-"],
                stdin=PIPE,
                stderr=PIPE,
            )
            outs, errs = process.communicate(input=key.encode("utf-8"))
            working = True
            while working:
                time.sleep(0.5)
                process.poll()
                rc = process.returncode
                if rc is None:
                    working = True
                else:
                    process = None
                    working = False
            if rc != 0:
                self.return_state = False
                self.return_message = _(
                    "Error {} while doing persistent partition: {}"
                ).format(rc, errs.decode("utf-8"))
                logging.error(self.return_message)
                self.finished.set()
                return
            self._progress = 50
            self.return_message = _("Persistent partition encrypted. Opening...")
            # cryptsetup open /dev/sdb3 crypt_sdb3
            if not self.udev_wait("creating encrypted partition"):
                return

            process = Popen(
                [
                    "cryptsetup",
                    "luksOpen",
                    target + "3",
                    "crypt_" + base_target,
                    "-d",
                    "-",
                ],
                stdin=PIPE,
                stderr=PIPE,
            )
            outs, errs = process.communicate(input=key.encode("utf-8"))
            working = True
            while working:
                time.sleep(0.5)
                process.poll()
                rc = process.returncode
                if rc is None:
                    working = True
                else:
                    process = None
                    working = False
            if rc != 0:
                self.return_state = False
                self.return_message = _(
                    "Error {} while doing persistent partition: {}"
                ).format(rc, errs.decode("utf-8"))
                logging.error(self.return_message)
                self.finished.set()
                return

            if not self.udev_wait(_("opening encrypted partition")):
                return
            self._progress = 60
            self.return_message = _("Persistent partition opened. Formatting...")
            # mkfs.ext4 -L mgalive-persist /dev/mapper/crypt_sdb3
            process = Popen(
                [
                    "mkfs.ext4",
                    "-q",
                    "-F",
                    "-L",
                    label,
                    "/dev/mapper/crypt_" + base_target,
                ],
                stderr=PIPE,
            )
            outs, errs = process.communicate()
            working = True
            while working:
                time.sleep(0.5)
                process.poll()
                rc = process.returncode
                if rc is None:
                    working = True
                else:
                    process = None
                    working = False
            if rc != 0:
                self.return_state = False
                self.return_message = _(
                    "Error {} while doing persistent partition: {}"
                ).format(rc, errs.decode("utf-8"))
                logging.error(self.return_message)
                self.finished.set()
                return
            self._progress = 90
            # cryptsetup close crypt_sdb3

            if not self.udev_wait(_("formatting encrypted partition")):
                return
            process = Popen(["cryptsetup", "luksClose", "crypt_" + base_target])
            outs, errs = process.communicate()
            working = True
            while working:
                time.sleep(0.5)
                process.poll()
                rc = process.returncode
                if rc is None:
                    working = True
                else:
                    process = None
                    working = False
            if rc != 0:
                self.return_state = False
                self.return_message = _(
                    "Error {} while doing persistent partition: {}"
                ).format(rc, errs.decode("utf-8"))
                logging.error(self.return_message)
            else:
                self.return_state = True
                self.return_message = _("Persistent partition done")
                logging.info(self.return_message)
            if not self.udev_wait(_("closing encrypted partition")):
                return
        self._progress = 100
        self.finished.set()
        logging.debug("Finished")

    def __init__(self):
        gettext.install("isodumper", localedir="/usr/share/locale")
        # Operation running
        self.operation = False