#!/usr/bin/env python
"""Helper script for running end-to-end tests."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import atexit
import os
import signal
import socket
import subprocess
import sys
import tempfile
import threading
import time

from absl import app
from absl import flags
import portpicker
import psutil
import requests

from grr_api_client import api
from grr_response_core import config
from grr_response_core.lib import config_lib
from grr_response_core.lib import package


class Error(Exception):
  """Module-specific base error class."""


class TCPPortTimeout(Error):
  """Raised when a TCP port didn't open in time."""


class ClientEnrollmentTimeout(Error):
  """Raised when a client does not enroll in time."""


flags.DEFINE_list(
    "tests", [],
    "(Optional) comma-separated list of tests to run (skipping all others). "
    "If this flag is not specified, all tests available for the platform "
    "will run.")

flags.DEFINE_list(
    "manual_tests", [],
    "A comma-separated list of extra tests to run (such tests are not run by "
    "default and have to be manually enabled with this flag).")

flags.DEFINE_string("mysql_database", "grr_test_db",
                    "MySQL database name to use.")

flags.DEFINE_string("mysql_username", None, "MySQL username to use.")

flags.DEFINE_string("mysql_password", None, "MySQL password to use.")

flags.DEFINE_string("logging_path", None,
                    "Base logging path for server components to use.")


def GetServerComponentArgs(config_path):
  """Returns a set of command line arguments for server components.

  Args:
    config_path: Path to a config path generated by
      self_contained_config_writer.

  Returns:
    An iterable with command line arguments to use.
  """

  primary_config_path = package.ResourcePath(
      "grr-response-core", "install_data/etc/grr-server.yaml")
  secondary_config_path = package.ResourcePath(
      "grr-response-test", "grr_response_test/test_data/grr_test.yaml")
  return [
      "--config",
      primary_config_path,
      "--secondary_configs",
      ",".join([secondary_config_path, config_path]),
      "-p",
      "Monitoring.http_port=%d" % portpicker.pick_unused_port(),
      "-p",
      "AdminUI.webauth_manager=NullWebAuthManager",
  ]


def StartComponent(main_package, args):
  """Starts a new process with a given component.

  This starts a Python interpreter with a "-m" argument followed by
  the main package name, thus effectively executing the main()
  function of a given package.

  Args:
    main_package: Main package path.
    args: An iterable with program arguments (not containing the program
      executable).

  Returns:
    Popen object corresponding to a started process.
  """
  popen_args = [sys.executable, "-m", main_package] + args
  print("Starting %s component: %s" % (main_package, " ".join(popen_args)))
  process = subprocess.Popen(popen_args)
  print("Component %s pid: %d" % (main_package, process.pid))
  return process


_PROCESS_CHECK_INTERVAL = 0.1


def DieIfSubProcessDies(processes):
  """Kills the process if any of given processes dies.

  This function is supposed to run in a background thread and monitor provided
  processes to ensure they don't die silently.

  Args:
    processes: An iterable with multiprocessing.Process instances.
  """
  while True:
    for p in processes:
      if p.poll() not in [None, 0]:
        # DieIfSubProcessDies runs in a background thread, raising an exception
        # will just kill the thread while what we want is to fail the whole
        # process.
        print("Subprocess %s died unexpectedly. Killing main process..." %
              p.pid)
        for kp in processes:
          try:
            os.kill(kp.pid, signal.SIGTERM)
          except OSError:
            pass
        # sys.exit only exits a thread when called from a thread.
        # Killing self with SIGTERM to ensure the process runs necessary
        # cleanups before exiting.
        os.kill(os.getpid(), signal.SIGTERM)
    time.sleep(_PROCESS_CHECK_INTERVAL)


_TCP_PORT_WAIT_TIMEOUT_SECS = 15


def WaitForTCPPort(port):
  """Waits for a given local TCP port to open.

  If the port in question does not open within ~10 seconds, main process gets
  killed.

  Args:
    port: An integer identifying the port.

  Raises:
    TCPPortTimeout: if the port doesn't open.
  """
  start_time = time.time()
  while time.time() - start_time < _TCP_PORT_WAIT_TIMEOUT_SECS:
    try:
      sock = socket.create_connection(("localhost", port))
      sock.close()
      return
    except socket.error:
      pass
    time.sleep(_PROCESS_CHECK_INTERVAL)

  raise TCPPortTimeout("TCP port %d didn't open." % port)


_CLIENT_ENROLLMENT_WAIT_TIMEOUT_SECS = 15
_CLIENT_ENROLLMENT_CHECK_INTERVAL = 1


def WaitForClientToEnroll(admin_ui_port):
  """Waits for an already started client to enroll.

  If the client doesn't enroll within ~100 seconds, main process gets killed.

  Args:
    admin_ui_port: AdminUI port to be used with API client library to check for
      an enrolled client.

  Returns:
    A string with an enrolled client's client id.

  Raises:
    ClientEnrollmentTimeout: if the client fails to enroll in time.
  """
  api_endpoint = "http://localhost:%d" % admin_ui_port

  start_time = time.time()
  while time.time() - start_time < _CLIENT_ENROLLMENT_WAIT_TIMEOUT_SECS * 10:
    try:
      api_client = api.InitHttp(api_endpoint=api_endpoint)
      clients = list(api_client.SearchClients(query="."))
    except requests.exceptions.ConnectionError:
      #      print("Connection error (%s), waiting..." % api_endpoint)
      time.sleep(_CLIENT_ENROLLMENT_CHECK_INTERVAL)
      continue

    if clients:
      return clients[0].client_id

    print("No clients enrolled, waiting...")
    time.sleep(_CLIENT_ENROLLMENT_CHECK_INTERVAL)

  raise ClientEnrollmentTimeout("Client didn't enroll.")


def GetRunEndToEndTestsArgs(client_id, server_config):
  """Returns arguments needed to configure run_end_to_end_tests process.

  Args:
    client_id: String with a client id pointing to an already running client.
    server_config: GRR configuration object with a server configuration.

  Returns:
    An iterable with command line arguments.
  """
  api_endpoint = "http://localhost:%d" % server_config["AdminUI.port"]
  args = [
      "--api_endpoint",
      api_endpoint,
      "--api_user",
      "admin",
      "--api_password",
      "admin",
      "--client_id",
      client_id,
      "--ignore_test_context",
      "True",
      "-p",
      "Client.binary_name=%s" % psutil.Process(os.getpid()).name(),
  ]
  if flags.FLAGS.tests:
    args += ["--whitelisted_tests", ",".join(flags.FLAGS.tests)]
  if flags.FLAGS.manual_tests:
    args += ["--manual_tests", ",".join(flags.FLAGS.manual_tests)]

  return args


def main(argv):
  del argv  # Unused.

  if flags.FLAGS.mysql_username is None:
    raise ValueError("--mysql_username has to be specified.")

  # Create 2 temporary files to contain server and client configuration files
  # that we're about to generate.
  #
  # TODO(user): migrate to TempFilePath as soon grr.test_lib is moved to
  # grr_response_test.
  fd, built_server_config_path = tempfile.mkstemp(".yaml")
  os.close(fd)
  print("Using temp server config path: %s" % built_server_config_path)
  fd, built_client_config_path = tempfile.mkstemp(".yaml")
  os.close(fd)
  print("Using temp client config path: %s" % built_client_config_path)

  def CleanUpConfigs():
    os.remove(built_server_config_path)
    os.remove(built_client_config_path)

  atexit.register(CleanUpConfigs)

  # Generate server and client configs.
  config_writer_flags = [
      "--dest_server_config_path",
      built_server_config_path,
      "--dest_client_config_path",
      built_client_config_path,
      "--config_mysql_database",
      flags.FLAGS.mysql_database,
  ]

  if flags.FLAGS.mysql_username is not None:
    config_writer_flags.extend(
        ["--config_mysql_username", flags.FLAGS.mysql_username])

  if flags.FLAGS.mysql_password is not None:
    config_writer_flags.extend(
        ["--config_mysql_password", flags.FLAGS.mysql_password])

  if flags.FLAGS.logging_path is not None:
    config_writer_flags.extend(
        ["--config_logging_path", flags.FLAGS.logging_path])

  p = StartComponent(
      "grr_response_test.lib.self_contained_config_writer",
      config_writer_flags)
  if p.wait() != 0:
    raise RuntimeError("ConfigWriter execution failed: {}".format(p.returncode))

  server_config = config_lib.LoadConfig(config.CONFIG.MakeNewConfig(),
                                        built_server_config_path)

  # Start all remaining server components.
  processes = [
      StartComponent(
          "grr_response_server.gui.admin_ui",
          GetServerComponentArgs(built_server_config_path)),
      StartComponent(
          "grr_response_server.bin.frontend",
          GetServerComponentArgs(built_server_config_path)),
      StartComponent(
          "grr_response_server.bin.worker",
          GetServerComponentArgs(built_server_config_path)),
      StartComponent(
          "grr_response_client.client",
          ["--config", built_client_config_path, "--verbose"]),
  ]

  # Start a background thread that kills the main process if one of the
  # subprocesses dies.
  t = threading.Thread(target=DieIfSubProcessDies, args=(processes,))
  t.daemon = True
  t.start()

  # Wait for the client to enroll and get its id.
  client_id = WaitForClientToEnroll(server_config["AdminUI.port"])
  print("Found client id: %s" % client_id)

  # Run the test suite against the enrolled client.
  p = StartComponent(
      "grr_response_test.run_end_to_end_tests",
      GetServerComponentArgs(built_server_config_path) +
      GetRunEndToEndTestsArgs(client_id, server_config))
  if p.wait() != 0:
    raise RuntimeError("RunEndToEndTests execution failed.")

  print("RunEndToEndTests execution succeeded.")


if __name__ == "__main__":
  app.run(main)
