# Copyright 2016 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import importlib
import logging

MOBLY_CONTROLLER_CONFIG_NAME = "Sniffer"


def create(configs):
  """Initializes the sniffer structures based on the JSON configuration. The
    expected keys are:

        * Type: A first-level type of sniffer. Planned to be 'local' for
            sniffers running on the local machine, or 'remote' for sniffers
            running remotely.
        * SubType: The specific sniffer type to be used.
        * Interface: The WLAN interface used to configure the sniffer.
        * BaseConfigs: A dictionary specifying baseline configurations of
            the sniffer. Configurations can be overridden when starting a
            capture. The keys must be one of the Sniffer.CONFIG_KEY_*
            values.
    """
  objs = []
  for c in configs:
    sniffer_type = c["Type"]
    sniffer_subtype = c["SubType"]
    interface = c["Interface"]
    base_configs = c["BaseConfigs"]
    module_name = "mobly.controllers.sniffer_lib.{}.{}".format(
        sniffer_type, sniffer_subtype)
    module = importlib.import_module(module_name)
    objs.append(
        module.Sniffer(interface,
                       logging.getLogger(),
                       base_configs=base_configs))
  return objs


def destroy(objs):
  """Destroys the sniffers and terminates any ongoing capture sessions.
    """
  for sniffer in objs:
    try:
      sniffer.stop_capture()
    except SnifferError:
      pass


class SnifferError(Exception):
  """This is the Exception class defined for all errors generated by
    Sniffer-related modules.
    """
  pass


class InvalidDataError(Exception):
  """This exception is thrown when invalid configuration data is passed
    to a method.
    """
  pass


class ExecutionError(SnifferError):
  """This exception is thrown when trying to configure the capture device
    or when trying to execute the capture operation.

    When this exception is seen, it is possible that the sniffer module is run
    without sudo (for local sniffers) or keys are out-of-date (for remote
    sniffers).
    """
  pass


class InvalidOperationError(SnifferError):
  """Certain methods may only be accessed when the instance upon which they
    are invoked is in a certain state. This indicates that the object is not
    in the correct state for a method to be called.
    """
  pass


class Sniffer:
  """This class defines an object representing a sniffer.

    The object defines the generic behavior of sniffers - irrespective of how
    they are implemented, or where they are located: on the local machine or on
    the remote machine.
    """

  CONFIG_KEY_CHANNEL = "channel"

  def __init__(self, interface, logger, base_configs=None):
    """The constructor for the Sniffer. It constructs a sniffer and
        configures it to be ready for capture.

        Args:
            interface: A string specifying the interface used to configure the
                sniffer.
            logger: Mobly logger object.
            base_configs: A dictionary containing baseline configurations of the
                sniffer. These can be overridden when staring a capture. The
                keys are specified by Sniffer.CONFIG_KEY_*.

        Returns:
            self: A configured sniffer.

        Raises:
            InvalidDataError: if the config_path is invalid.
            NoPermissionError: if an error occurs while configuring the
                sniffer.
        """
    raise NotImplementedError("Base class should not be called directly!")

  def get_descriptor(self):
    """This function returns a string describing the sniffer. The specific
        string (and its format) is up to each derived sniffer type.

        Returns:
            A string describing the sniffer.
        """
    raise NotImplementedError("Base class should not be called directly!")

  def get_type(self):
    """This function returns the type of the sniffer.

        Returns:
            The type (string) of the sniffer. Corresponds to the 'Type' key of
            the sniffer configuration.
        """
    raise NotImplementedError("Base class should not be called directly!")

  def get_subtype(self):
    """This function returns the sub-type of the sniffer.

        Returns:
            The sub-type (string) of the sniffer. Corresponds to the 'SubType'
            key of the sniffer configuration.
        """
    raise NotImplementedError("Base class should not be called directly!")

  def get_interface(self):
    """This function returns The interface used to configure the sniffer,
        e.g. 'wlan0'.

        Returns:
            The interface (string) used to configure the sniffer. Corresponds to
            the 'Interface' key of the sniffer configuration.
        """
    raise NotImplementedError("Base class should not be called directly!")

  def get_capture_file(self):
    """The sniffer places a capture in the logger directory. This function
        enables the caller to obtain the path of that capture.

        Returns:
            The full path of the current or last capture.
        """
    raise NotImplementedError("Base class should not be called directly!")

  def start_capture(self,
                    override_configs=None,
                    additional_args=None,
                    duration=None,
                    packet_count=None):
    """This function starts a capture which is saved to the specified file
        path.

        Depending on the type/subtype and configuration of the sniffer the
        capture may terminate on its own or may require an explicit call to the
        stop_capture() function.

        This is a non-blocking function so a terminating function must be
        called either explicitly or implicitly:

            * Explicitly: call either stop_capture() or wait_for_capture()
            * Implicitly: use with a with clause. The wait_for_capture()
                function will be called if a duration is specified (i.e.
                is not None), otherwise a stop_capture() will be called.

        The capture is saved to a file in the log path of the logger. Use
        the get_capture_file() to get the full path to the current or most
        recent capture.

        Args:
            override_configs: A dictionary which is combined with the
                base_configs ("BaseConfigs" in the sniffer configuration). The
                keys (specified by Sniffer.CONFIG_KEY_*) determine the
                configuration of the sniffer for this specific capture.
            additional_args: A string specifying additional raw
                command-line arguments to pass to the underlying sniffer. The
                interpretation of these flags is sniffer-dependent.
            duration: An integer specifying the number of seconds over which to
                capture packets. The sniffer will be terminated after this
                duration. Used in implicit mode when using a 'with' clause. In
                explicit control cases may have to be performed using a
                sleep+stop or as the timeout argument to the wait function.
            packet_count: An integer specifying the number of packets to capture
                before terminating. Should be used with duration to guarantee
                that capture terminates at some point (even if did not capture
                the specified number of packets).

        Returns:
            An ActiveCaptureContext process which can be used with a 'with'
            clause.

        Raises:
            InvalidDataError: for invalid configurations
            NoPermissionError: if an error occurs while configuring and running
                the sniffer.
        """
    raise NotImplementedError("Base class should not be called directly!")

  def stop_capture(self):
    """This function stops a capture and guarantees that the capture is
        saved to the capture file configured during the start_capture() method.
        Depending on the type of the sniffer the file may previously contain
        partial results (e.g. for a local sniffer) or may not exist until the
        stop_capture() method is executed (e.g. for a remote sniffer).

        Depending on the type/subtype and configuration of the sniffer the
        capture may terminate on its own without requiring a call to this
        function. In such a case it is still necessary to call either this
        function or the wait_for_capture() function to make sure that the
        capture file is moved to the correct location.

        Raises:
            NoPermissionError: No permission when trying to stop a capture
                and save the capture file.
        """
    raise NotImplementedError("Base class should not be called directly!")

  def wait_for_capture(self, timeout=None):
    """This function waits for a capture to terminate and guarantees that
        the capture is saved to the capture file configured during the
        start_capture() method. Depending on the type of the sniffer the file
        may previously contain partial results (e.g. for a local sniffer) or
        may not exist until the stop_capture() method is executed (e.g. for a
        remote sniffer).

        Depending on the type/subtype and configuration of the sniffer the
        capture may terminate on its own without requiring a call to this
        function. In such a case it is still necessary to call either this
        function or the stop_capture() function to make sure that the capture
        file is moved to the correct location.

        Args:
            timeout: An integer specifying the number of seconds to wait for
                the capture to terminate on its own. On expiration of the
                timeout the sniffer is stopped explicitly using the
                stop_capture() function.

        Raises:
            NoPermissionError: No permission when trying to stop a capture and
                save the capture file.
        """
    raise NotImplementedError("Base class should not be called directly!")


class ActiveCaptureContext:
  """This class defines an object representing an active sniffer capture.

    The object is returned by a Sniffer.start_capture() command and terminates
    the capture when the 'with' clause exits. It is syntactic sugar for
    try/finally.
    """

  _sniffer = None
  _timeout = None

  def __init__(self, sniffer, timeout=None):
    self._sniffer = sniffer
    self._timeout = timeout

  def __enter__(self):
    pass

  def __exit__(self, type, value, traceback):
    if self._sniffer is not None:
      if self._timeout is None:
        self._sniffer.stop_capture()
      else:
        self._sniffer.wait_for_capture(self._timeout)
    self._sniffer = None