# -*- coding: utf-8 -*-
"""Creates an analysis VM and attaches GCP disks to it for analysis."""

import time
from typing import Callable, Dict

from libcloudforensics import errors as lcf_errors
from libcloudforensics.providers.gcp import forensics as gcp_forensics
from libcloudforensics.providers.gcp.internal import common
from libcloudforensics.providers.gcp.internal import project as gcp_project
from libcloudforensics.providers.gcp.internal import compute

from dftimewolf.lib import module
from dftimewolf.lib.containers import containers
from dftimewolf.lib.modules import manager as modules_manager
from dftimewolf.lib import cache
from dftimewolf.lib import telemetry
from dftimewolf.lib.containers import manager as container_manager


class GCEForensicsVM(module.BaseModule):
  """Google Cloud Forensics VM creator.

  Attributes:
    project (gcp.GoogleCloudProject): Project in which to create the VM.
    analysis_vm (gcp.GoogleComputeInstance): Analysis VM on
         which the disk copy will be attached.
    incident_id (str): Incident identifier used to name the analysis VM.
    boot_disk_size (Optional[float]): Optional. Size of the analysis VM boot
        disk (in GB).
    boot_disk_type (Optional[str]): Optional. Disk type to use.
    cpu_cores (Optional[int]): Optional. Number of CPU cores to create the VM
        with.
    image_project (Optional[str]): Optional. Name of the project where the
        analysis VM image is hosted.
    image_family (Optional[str]): Optional. Name of the image to use to
        create the analysis VM.
    create_analysis_vm: Legacy option. False to skip this module entirely.
  """

  _ANALYSIS_VM_CONTAINER_ATTRIBUTE_NAME = 'Analysis VM'
  _ANALYSIS_VM_CONTAINER_ATTRIBUTE_TYPE = 'text'

  def __init__(self,
               name: str,
               container_manager_: container_manager.ContainerManager,
               cache_: cache.DFTWCache,
               telemetry_: telemetry.BaseTelemetry,
               publish_message_callback: Callable[[str, str, bool], None]):
    """Initializes a Google Cloud Platform (GCP) collector.

    Args:
      name: The modules runtime name.
      container_manager_: A common container manager object.
      cache_: A common DFTWCache object.
      telemetry_: A common telemetry collector object.
      publish_message_callback: A callback to send modules messages to.
    """
    super().__init__(name=name,
                     cache_=cache_,
                     container_manager_=container_manager_,
                     telemetry_=telemetry_,
                     publish_message_callback=publish_message_callback)

    self.project = None  # type: gcp_project.GoogleCloudProject
    self.analysis_vm = None  # type: compute.GoogleComputeInstance
    self.analysis_vm_name = ''
    self.incident_id = str()
    self.boot_disk_size = 0.0
    self.boot_disk_type = str()
    self.cpu_cores = 0
    self.image_project = str()
    self.image_family = str()
    self._gcp_label = {}  # type: Dict[str, str]
    self.create_analysis_vm = bool()

  # pylint: disable=arguments-differ,too-many-arguments
  def SetUp(self,
            project_name: str,
            incident_id: str,
            zone: str,
            boot_disk_size: float,
            boot_disk_type: str,
            cpu_cores: int,
            image_project: str,
            image_family: str,
            create_analysis_vm: bool,
            analysis_vm_name: str) -> None:
    """Sets up a GCE Forensics VM processor.

    Args:
      project_name: Optional. name of the project that contains the analysis VM.
      zone: Optional. GCP zone in which new resources should be created.
      boot_disk_size: Size of the analysis VM boot disk (in GB). Default is 50.
      boot_disk_type: Disk type to use.
      cpu_cores: Number of CPU cores to create the VM with.
      image_project: Name of the project where the analysis VM image is hosted.
      image_family: Name of the image to use to create the analysis VM.
      create_analysis_vm: Legacy option. False to skip this module entirely.
      analysis_vm_name: A name to use for the created analysis vm. Actual name
          will also be suffixed with the incident ID (if the ID is not already
          in the name.)
    """
    self.create_analysis_vm = create_analysis_vm
    if not self.create_analysis_vm:
      self.logger.warning('Skipping SetUp for Forensics VM creation.')
      return

    self.project = gcp_project.GoogleCloudProject(
        project_name, default_zone=zone)

    if incident_id:
      self.incident_id = incident_id
      self._gcp_label = {'incident_id': self.incident_id}

    self.boot_disk_size = boot_disk_size
    self.boot_disk_type = boot_disk_type
    self.cpu_cores = cpu_cores
    self.image_project = image_project
    self.image_family = image_family

    self.analysis_vm_name = analysis_vm_name or 'gcp-forensics-vm'
    if self.incident_id:
      if self.incident_id not in self.analysis_vm_name:
        self.analysis_vm_name = f'{self.analysis_vm_name}-{self.incident_id}'
    else:
      self.analysis_vm_name = common.GenerateUniqueInstanceName(
          self.analysis_vm_name,
          common.COMPUTE_NAME_LIMIT)

  def Process(self) -> None:
    """Launches the analysis VM."""
    if not self.create_analysis_vm:
      self.logger.warning('Skipping Process for Forensics VM creation.')
      return

    self.PublishMessage(f'Your analysis VM will be: {self.analysis_vm_name}')
    self.logger.info("Complimentary gcloud command:")
    self.logger.info(
      f"gcloud compute ssh --project {self.project.project_id} "
      f"{self.analysis_vm_name} --zone {self.project.default_zone}"
    )
    self.StoreContainer(
        containers.TicketAttribute(
            name=self._ANALYSIS_VM_CONTAINER_ATTRIBUTE_NAME,
            type_=self._ANALYSIS_VM_CONTAINER_ATTRIBUTE_TYPE,
            value=self.analysis_vm_name))
    try:
      # pylint: disable=too-many-function-args
      # pylint: disable=redundant-keyword-arg
      self.analysis_vm, created = gcp_forensics.StartAnalysisVm(
          self.project.project_id,
          self.analysis_vm_name,
          self.project.default_zone,
          self.boot_disk_size,
          self.boot_disk_type,
          int(self.cpu_cores),
          image_project=self.image_project,
          image_family=self.image_family)
    except lcf_errors.ResourceCreationError as exception:
      self.logger.error(f'Could not create VM: {exception}')
      self.ModuleError(str(exception), critical=True)
    if not created:
      self.logger.debug(f"Instance {self.analysis_vm_name} exists: reusing.")
    if self._gcp_label:
      self.analysis_vm.AddLabels(self._gcp_label)
      self.analysis_vm.GetBootDisk().AddLabels(self._gcp_label)
    self.StoreContainer(containers.ForensicsVM(
        name=self.analysis_vm.name,
        evidence_disk=None,
        platform='gcp'))

    disks = self.GetContainers(containers.GCEDisk)

    # Sleep until status is RUNNING before attaching disks. Possible values:
    # https://cloud.google.com/compute/docs/reference/rest/v1/instances/get
    if disks and self.analysis_vm.GetPowerState() != 'RUNNING':
      self.logger.debug(
        "Pausing 10 seconds to allow OS to boot before attaching"
        " evidence disks"
      )
      time.sleep(10)
    # In Running boot might be ongoing, thus the additinal wait
    # https://cloud.google.com/compute/docs/instances/instance-life-cycle
    time.sleep(20)

    for d in disks:
      if d.project != self.project.project_id:
        continue
      try:
        self.logger.info(f'Attaching {d.name} to {self.analysis_vm_name}')
        self.analysis_vm.AttachDisk(compute.GoogleComputeDisk(
            self.project.project_id,
            self.project.default_zone,
            d.name))
      except RuntimeError as error:
        if 'RESOURCE_IN_USE_BY_ANOTHER_RESOURCE' in str(error):
          self.logger.warning(
              f'Attaching {d.name} to {self.analysis_vm_name} failed, as it is '
              'already attached to another instance.')
        else:
          raise error


modules_manager.ModulesManager.RegisterModule(GCEForensicsVM)
