#!/usr/bin/python3
"""
Copyright (C) 2023  Michael Ablassmeier <abi@grinser.de>

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 3 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, see <https://www.gnu.org/licenses/>.
"""
import os
import json
import logging
from argparse import Namespace
from typing import List, Union
import libvirt

from libvirtnbdbackup import output
from libvirtnbdbackup.virt.client import DomainDisk
from libvirtnbdbackup import common as lib
from libvirtnbdbackup.qemu import util as qemu
from libvirtnbdbackup.qemu.exceptions import ProcessError
from libvirtnbdbackup.ssh.exceptions import sshError
from libvirtnbdbackup.output.exceptions import OutputException
from libvirtnbdbackup.output.target.directory import Directory
from libvirtnbdbackup.output.target.zip import Zip
from libvirtnbdbackup.common import safeInfo
from libvirtnbdbackup.virt import guest


log = logging.getLogger()


def backupChecksum(fileStream, targetFile):
    """Save the calculated adler32 checksum, it can be verified
    by virtnbdbrestore's verify function.'"""
    checksum = fileStream.checksum()
    safeInfo("Checksum for file: [%s]:[%s]", targetFile, checksum)
    chksumfile = f"{targetFile}.chksum"
    safeInfo("Saving checksum to: [%s]", chksumfile)
    with output.openfile(chksumfile, "w") as cf:
        cf.write(f"{checksum}")


def backupConfig(args: Namespace, vmConfig: str) -> Union[str, None]:
    """Save domain XML config file"""
    configFile = os.path.join(args.output, f"vmconfig.{lib.getIdent(args)}.xml")
    log.info("Saving VM config to: [%s]", configFile)
    try:
        with output.openfile(configFile, "wb") as fh:
            fh.write(vmConfig.encode())
        return configFile
    except OutputException as e:
        log.error("Failed to save VM config: [%s]", e)
        return None


def backupDiskInfo(args: Namespace, disk: DomainDisk):
    """Save information about qcow image, used to reconstruct
    the qemu image with the same settings during restore"""
    try:
        info = qemu.util("").info(disk.path, args.sshClient)
    except (
        ProcessError,
        sshError,
    ) as errmsg:
        log.warning("Failed to read qcow image info: [%s]", errmsg)
        return

    configFile = f"{args.output}/{disk.target}.{lib.getIdent(args)}.qcow.json"
    try:
        with output.openfile(configFile, "wb") as fh:
            fh.write(info.out.encode())
        log.info("Saved qcow image config to: [%s]", configFile)
        if args.stdout is True:
            args.diskInfo.append(configFile)
    except OutputException as e:
        log.warning("Failed to save qcow image config: [%s]", e)


def backupBootConfig(args: Namespace) -> None:
    """Save domain uefi/nvram/kernel and loader if configured."""
    for setting, val in args.info.items():
        if args.level != "copy":
            tFile = os.path.join(
                args.output, f"{os.path.basename(val)}.{lib.getIdent(args)}"
            )
        else:
            tFile = os.path.join(args.output, f"{os.path.basename(val)}")
        log.info("Save additional boot config [%s] to: [%s]", setting, tFile)
        lib.copy(args, val, tFile)
        args.info[setting] = tFile


def backupAutoStart(args: Namespace) -> None:
    """Save information if virtual machine was marked
    for autostart during system boot"""
    log.info("Autostart setting configured for virtual machine.")
    autoStartFile = os.path.join(args.output, f"autostart.{lib.getIdent(args)}")
    try:
        with output.openfile(autoStartFile, "wb") as fh:
            fh.write(b"True")
    except OutputException as e:
        log.warning("Failed to save autostart information: [%s]", e)


def backupGuestInfo(args: Namespace) -> None:
    """Save OS related information"""
    osInfoFile = os.path.join(args.output, f"osinfo.{lib.getIdent(args)}")
    try:
        with output.openfile(osInfoFile, "w") as fh:
            fh.write(json.dumps(args.guestInfo, indent=4))
        log.info("Saved guest related osinfo to [%s]", osInfoFile)
    except OutputException as e:
        log.warning("Failed to save osinfo data: [%s]", e)


def backupBitlockerRecoveryKey(args: Namespace, domObj: libvirt.virDomain) -> None:
    """Save bitlocker recovery keys"""
    try:
        bde = guest.Exec(domObj, "manage-bde.exe", ["-status"])
        log.info("Bitlocker tools detected, attempting to backup recovery keys.")
        log.debug(bde)
    except libvirt.libvirtError:
        log.info("System does not appear to have bitlocker tools installed, skipping.")

    for i in range(0, args.guestInfo["fs.count"]):
        vol = args.guestInfo.get(f"fs.{i}.mountpoint", None)
        if not vol:
            continue
        if not ":" in vol:
            log.info("Skipping volume: [%s]", vol)
            continue

        vol = vol.replace("\\", "")
        log.info("Get bitlocker recovery key for volume: [%s]", vol)
        try:
            protectors = guest.Exec(
                domObj, "manage-bde.exe", ["-protectors", "-get", vol]
            )
        except RuntimeError:
            log.warning(
                "Unable to extract recovery key for volume [%s], see debug log for error details.",
                vol,
            )
            continue
        except (TimeoutError, libvirt.libvirtError) as e:
            log.warning("Unable to pull recovery keys for [%s]: %s", vol, e)
            continue

        keyFile = os.path.join(
            args.output,
            f"bitlocker.recovery.key.{vol.replace(':','')}.{lib.getIdent(args)}",
        )
        try:
            with output.openfile(keyFile, "w") as fh:
                fh.write(protectors)
            log.info("Saved Bitlocker recovery key to [%s]", keyFile)
        except OutputException as e:
            log.warning("Failed to save recovery key: [%s]", e)


def saveFiles(
    args: Namespace,
    vmConfig: str,
    disks: List[DomainDisk],
    fileStream: Union[Directory, Zip],
    logFile: str,
):
    """Save additional files such as virtual machine configuration
    and UEFI / kernel images"""
    configFile = backupConfig(args, vmConfig)

    backupBootConfig(args)
    for disk in disks:
        if disk.format.startswith("qcow"):
            backupDiskInfo(args, disk)
    if args.stdout is True:
        addFiles(args, configFile, fileStream, logFile)


def addFiles(args: Namespace, configFile: Union[str, None], zipStream, logFile: str):
    """Add backup log and other files to zip archive"""
    if configFile is not None:
        log.info("Adding vm config to zipfile")
        zipStream.zipStream.write(configFile, configFile)
    if args.level in ("full", "inc"):
        log.info("Adding checkpoint info to zipfile")
        zipStream.zipStream.write(args.cpt.file, args.cpt.file)
        for dirname, _, files in os.walk(args.checkpointdir):
            zipStream.zipStream.write(dirname)
            for filename in files:
                zipStream.zipStream.write(os.path.join(dirname, filename))

    for setting, val in args.info.items():
        log.info("Adding additional [%s] setting file [%s] to zipfile", setting, val)
        zipStream.zipStream.write(val, os.path.basename(val))

    for diskInfo in args.diskInfo:
        log.info("Adding QCOW image format file [%s] to zipfile", diskInfo)
        zipStream.zipStream.write(diskInfo, os.path.basename(diskInfo))

    log.info("Adding backup log [%s] to zipfile", logFile)
    zipStream.zipStream.write(logFile, logFile)
