'''
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements.  See the NOTICE file
distributed with this work for additional information
regarding copyright ownership.  The ASF licenses this file
to you 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.
'''

from stacks.utils.RMFTestCase import *
import bootstrap
import time
from ambari_commons import subprocess32
import os
import logging
import tempfile
import pprint

from ambari_commons.os_check import OSCheck
from bootstrap import PBootstrap, Bootstrap, BootstrapDefault, SharedState, HostLog, SCP, SSH
from unittest import TestCase
from ambari_commons.subprocess32 import Popen
from bootstrap import AMBARI_PASSPHRASE_VAR_NAME
from mock.mock import MagicMock, call
from mock.mock import patch
from mock.mock import create_autospec
from only_for_platform import not_for_platform, os_distro_value, PLATFORM_WINDOWS

@not_for_platform(PLATFORM_WINDOWS)
class TestBootstrap:#(TestCase):

  def setUp(self):
    logging.basicConfig(level=logging.ERROR)


  def test_getRemoteName(self):
    shared_state = SharedState("root", "123", "sshkey_file", "scriptDir", "bootdir",
                      "setupAgentFile", "ambariServer", "centos6", None, "8440", "root")
    res = bootstrap_obj = Bootstrap("hostname", shared_state)
    utime1 = 1234
    utime2 = 12345
    bootstrap_obj.getUtime = MagicMock(return_value=utime1)
    remote1 = bootstrap_obj.getRemoteName("/tmp/setupAgent.sh")
    self.assertEquals(remote1, "/tmp/setupAgent{0}.sh".format(utime1))

    bootstrap_obj.getUtime.return_value=utime2
    remote1 = bootstrap_obj.getRemoteName("/tmp/setupAgent.sh")
    self.assertEquals(remote1, "/tmp/setupAgent{0}.sh".format(utime1))

    remote2 = bootstrap_obj.getRemoteName("/tmp/host_pass")
    self.assertEquals(remote2, "/tmp/host_pass{0}".format(utime2))


  # TODO: Test bootstrap timeout

  # TODO: test_return_error_message_for_missing_sudo_package

  def test_getAmbariPort(self):
    shared_state = SharedState("root", "123", "sshkey_file", "scriptDir", "bootdir",
                               "setupAgentFile", "ambariServer", "centos6",
                               None, "8440", "root")
    bootstrap_obj = Bootstrap("hostname", shared_state)
    self.assertEquals(bootstrap_obj.getAmbariPort(),"8440")
    shared_state.server_port = None
    bootstrap_obj = Bootstrap("hostname", shared_state)
    self.assertEquals(bootstrap_obj.getAmbariPort(),"null")


  @patch.object(subprocess32, "Popen")
  @patch("sys.stderr")
  @patch("sys.exit")
  @patch.object(PBootstrap, "run")
  @patch("os.path.dirname")
  @patch("os.path.realpath")
  def test_bootstrap_main(self, dirname_mock, realpath_mock, run_mock, exit_mock, stderr_mock, subprocess32_Popen_mock):
    bootstrap.main(["bootstrap.py", "hostname,hostname2", "/tmp/bootstrap", "root", "123", "sshkey_file", "setupAgent.py", "ambariServer", \
                    "centos6", "1.1.1", "8440", "root", "passwordfile"])
    self.assertTrue(run_mock.called)
    run_mock.reset_mock()
    bootstrap.main(["bootstrap.py", "hostname,hostname2", "/tmp/bootstrap", "root", "123", "sshkey_file", "setupAgent.py", "ambariServer", \
                    "centos6", "1.1.1", "8440", "root", None])
    self.assertTrue(run_mock.called)
    run_mock.reset_mock()
    def side_effect(retcode):
      raise Exception(retcode, "sys.exit")
    exit_mock.side_effect = side_effect
    try:
      bootstrap.main(["bootstrap.py","hostname,hostname2", "/tmp/bootstrap"])
      self.fail("sys.exit(2)")
    except Exception:
    # Expected
      pass
    self.assertTrue(exit_mock.called)


  @patch("os.environ")
  def test_getRunSetupWithPasswordCommand(self, environ_mock):
    shared_state = SharedState("root", "123", "sshkey_file", "scriptDir", "bootdir",
                               "setupAgentFile", "ambariServer", "centos6",
                               None, "8440", "root")
    environ_mock.__getitem__.return_value = "TEST_PASSPHRASE"
    bootstrap_obj = Bootstrap("hostname", shared_state)
    utime = 1234
    bootstrap_obj.getUtime = MagicMock(return_value=utime)
    ret = bootstrap_obj.getRunSetupWithPasswordCommand("hostname")
    expected = "/var/lib/ambari-agent/tmp/ambari-sudo.sh -S python /var/lib/ambari-agent/tmp/setupAgent{0}.py hostname TEST_PASSPHRASE " \
               "ambariServer root  8440 < /var/lib/ambari-agent/tmp/host_pass{0}".format(utime)
    self.assertEquals(ret, expected)


  def test_generateRandomFileName(self):
    shared_state = SharedState("root", "123", "sshkey_file", "scriptDir", "bootdir",
                               "setupAgentFile", "ambariServer", "centos6",
                               None, "8440", "root")
    bootstrap_obj = Bootstrap("hostname", shared_state)
    self.assertTrue(bootstrap_obj.generateRandomFileName(None) == bootstrap_obj.getUtime())



  @patch.object(OSCheck, "is_redhat_family")
  @patch.object(OSCheck, "is_suse_family")
  def test_getRepoDir(self, is_suse_family, is_redhat_family):
    shared_state = SharedState("root", "123", "sshkey_file", "scriptDir", "bootdir",
                               "setupAgentFile", "ambariServer", "centos6",
                               None, "8440", "root")
    bootstrap_obj = Bootstrap("hostname", shared_state)
    # Suse
    is_redhat_family.return_value = False
    is_suse_family.return_value = True
    res = bootstrap_obj.getRepoDir()
    self.assertEquals(res, "/etc/zypp/repos.d")
    # non-Suse
    is_suse_family.return_value = False
    is_redhat_family.return_value = True
    res = bootstrap_obj.getRepoDir()
    self.assertEquals(res, "/etc/yum.repos.d")

  def test_getSetupScript(self):
    shared_state = SharedState("root", "123", "sshkey_file", "scriptDir", "bootdir",
                               "setupAgentFile", "ambariServer", "centos6",
                               None, "8440", "root")
    bootstrap_obj = Bootstrap("hostname", shared_state)
    self.assertEquals(bootstrap_obj.shared_state.script_dir, "scriptDir")


  def test_run_setup_agent_command_ends_with_project_version(self):
    os.environ[AMBARI_PASSPHRASE_VAR_NAME] = ""
    version = "1.1.1"
    shared_state = SharedState("root", "123", "sshkey_file", "scriptDir", "bootdir",
                               "setupAgentFile", "ambariServer", "centos6",
                               version, "8440", "root")
    bootstrap_obj = Bootstrap("hostname", shared_state)
    runSetupCommand = bootstrap_obj.getRunSetupCommand("hostname")
    self.assertTrue(runSetupCommand.endswith(version + " 8440"))


  def test_agent_setup_command_without_project_version(self):
    os.environ[AMBARI_PASSPHRASE_VAR_NAME] = ""
    version = None
    shared_state = SharedState("root", "123", "sshkey_file", "scriptDir", "bootdir",
                               "setupAgentFile", "ambariServer", "centos6",
                               version, "8440", "root")
    bootstrap_obj = Bootstrap("hostname", shared_state)
    runSetupCommand = bootstrap_obj.getRunSetupCommand("hostname")
    self.assertTrue(runSetupCommand.endswith(" 8440"))


  # TODO: test_os_check_fail_fails_bootstrap_execution


  def test_host_log(self):
    tmp_file, tmp_filename  = tempfile.mkstemp()
    dummy_log = HostLog(tmp_filename)
    # First write to log
    dummy_log.write("a\nb\nc")
    # Read it
    with open(tmp_filename) as f:
      s = f.read()
      etalon = "a\nb\nc\n"
      self.assertEquals(s, etalon)
    # Next write
    dummy_log.write("Yet another string")
    # Read it
    with open(tmp_filename) as f:
      s = f.read()
      etalon = "a\nb\nc\nYet another string\n"
      self.assertEquals(s, etalon)
    # Should not append line end if it already exists
    dummy_log.write("line break->\n")
    # Read it
    with open(tmp_filename) as f:
      s = f.read()
      etalon = "a\nb\nc\nYet another string\nline break->\n"
      self.assertEquals(s, etalon)
    # Cleanup
    os.unlink(tmp_filename)


  @patch.object(subprocess32, "Popen")
  def test_SCP(self, popenMock):
    params = SharedState("root", "123", "sshkey_file", "scriptDir", "bootdir",
                                  "setupAgentFile", "ambariServer", "centos6",
                                  "1.2.1", "8440", "root")
    host_log_mock = MagicMock()
    log = {'text': ""}
    def write_side_effect(text):
      log['text'] = log['text'] + text

    host_log_mock.write.side_effect = write_side_effect
    scp = SCP(params.user, params.sshPort, params.sshkey_file, "dummy-host", "src/file",
              "dst/file", params.bootdir, host_log_mock)
    log_sample = "log_sample"
    error_sample = "error_sample"
    # Successful run
    process = MagicMock()
    popenMock.return_value = process
    process.communicate.return_value = (log_sample, error_sample)
    process.returncode = 0

    retcode = scp.run()

    self.assertTrue(popenMock.called)
    self.assertTrue(log_sample in log['text'])
    self.assertTrue(error_sample in log['text'])
    command_str = str(popenMock.call_args[0][0])
    self.assertEquals(command_str, "['scp', '-r', '-o', 'ConnectTimeout=60', '-o', "
        "'BatchMode=yes', '-o', 'StrictHostKeyChecking=no', '-P', '123', '-i', 'sshkey_file',"
        " 'src/file', 'root@dummy-host:dst/file']")
    self.assertEqual(retcode["exitstatus"], 0)

    log['text'] = ""
    #unsuccessfull run
    process.returncode = 1

    retcode = scp.run()

    self.assertTrue(log_sample in log['text'])
    self.assertTrue(error_sample in log['text'])
    self.assertEqual(retcode["exitstatus"], 1)


  @patch.object(subprocess32, "Popen")
  def test_SSH(self, popenMock):
    params = SharedState("root", "123", "sshkey_file", "scriptDir", "bootdir",
                                  "setupAgentFile", "ambariServer", "centos6",
                                  "1.2.1", "8440", "root")
    host_log_mock = MagicMock()
    log = {'text': ""}
    def write_side_effect(text):
      log['text'] = log['text'] + text

    host_log_mock.write.side_effect = write_side_effect
    ssh = SSH(params.user, params.sshPort, params.sshkey_file, "dummy-host", "dummy-command",
              params.bootdir, host_log_mock)
    log_sample = "log_sample"
    error_sample = "error_sample"
    # Successful run
    process = MagicMock()
    popenMock.return_value = process
    process.communicate.return_value = (log_sample, error_sample)
    process.returncode = 0

    retcode = ssh.run()

    self.assertTrue(popenMock.called)
    self.assertTrue(log_sample in log['text'])
    self.assertTrue(error_sample in log['text'])
    command_str = str(popenMock.call_args[0][0])
    self.assertEquals(command_str, "['ssh', '-o', 'ConnectTimeOut=60', '-o', "
            "'StrictHostKeyChecking=no', '-o', 'BatchMode=yes', '-tt', '-i', "
            "'sshkey_file', '-p', '123', 'root@dummy-host', 'dummy-command']")
    self.assertEqual(retcode["exitstatus"], 0)

    log['text'] = ""
    #unsuccessfull run
    process.returncode = 1

    retcode = ssh.run()

    self.assertTrue(log_sample in log['text'])
    self.assertTrue(error_sample in log['text'])
    self.assertEqual(retcode["exitstatus"], 1)

    log['text'] = ""
    # unsuccessful run with error message
    process.returncode = 1

    dummy_error_message = "dummy_error_message"
    ssh = SSH(params.user, params.sshPort, params.sshkey_file, "dummy-host", "dummy-command",
              params.bootdir, host_log_mock, errorMessage= dummy_error_message)
    retcode = ssh.run()

    self.assertTrue(log_sample in log['text'])
    self.assertTrue(error_sample in log['text'])
    self.assertTrue(dummy_error_message in log['text'])
    self.assertEqual(retcode["exitstatus"], 1)


  def test_getOsCheckScript(self):
    shared_state = SharedState("root", "123", "sshkey_file", "scriptDir", "bootdir",
                               "setupAgentFile", "ambariServer", "centos6",
                               None, "8440", "root")
    bootstrap_obj = Bootstrap("hostname", shared_state)
    ocs = bootstrap_obj.getOsCheckScript()
    self.assertEquals(ocs, "scriptDir/os_check_type.py")


  @patch.object(BootstrapDefault, "getRemoteName")
  def test_getOsCheckScriptRemoteLocation(self, getRemoteName_mock):
    shared_state = SharedState("root", "123", "sshkey_file", "scriptDir", "bootdir",
                               "setupAgentFile", "ambariServer", "centos6",
                               None, "8440", "root")
    bootstrap_obj = Bootstrap("hostname", shared_state)
    v = "/tmp/os_check_type1374259902.py"
    getRemoteName_mock.return_value = v
    ocs = bootstrap_obj.getOsCheckScriptRemoteLocation()
    self.assertEquals(ocs, v)


  @patch.object(BootstrapDefault, "is_suse")
  def test_getRepoFile(self, is_suse_mock):
    shared_state = SharedState("root", "123", "sshkey_file", "scriptDir", "bootdir",
                               "setupAgentFile", "ambariServer", "centos6",
                               None, "8440", "root")
    bootstrap_obj = Bootstrap("hostname", shared_state)
    is_suse_mock.return_value = False
    rf = bootstrap_obj.getRepoFile()
    self.assertEquals(rf, "/etc/yum.repos.d/ambari.repo")


  @patch.object(SSH, "__init__")
  @patch.object(SSH, "run")
  @patch.object(HostLog, "write")
  def test_createTargetDir(self, write_mock, run_mock,
                            init_mock):
    shared_state = SharedState("root", "123", "sshkey_file", "scriptDir", "bootdir",
                               "setupAgentFile", "ambariServer", "centos6",
                               None, "8440", "root")
    bootstrap_obj = Bootstrap("hostname", shared_state)
    expected = 42
    init_mock.return_value = None
    run_mock.return_value = expected
    res = bootstrap_obj.createTargetDir()
    self.assertEquals(res, expected)
    command = str(init_mock.call_args[0][4])
    self.assertEqual(command,
                     "SUDO=$([ \"$EUID\" -eq 0 ] && echo || echo sudo) ; $SUDO mkdir -p /var/lib/ambari-agent/tmp ; "
                     "$SUDO chown -R root /var/lib/ambari-agent/tmp ; "
                     "$SUDO chmod 755 /var/lib/ambari-agent ; "
                     "$SUDO chmod 755 /var/lib/ambari-agent/data ; "
                     "$SUDO chmod 1777 /var/lib/ambari-agent/tmp")

  @patch.object(BootstrapDefault, "getOsCheckScript")
  @patch.object(BootstrapDefault, "getOsCheckScriptRemoteLocation")
  @patch.object(SCP, "__init__")
  @patch.object(SCP, "run")
  @patch.object(HostLog, "write")
  def test_copyOsCheckScript(self, write_mock, run_mock, init_mock,
                    getOsCheckScriptRemoteLocation_mock, getOsCheckScript_mock):
    shared_state = SharedState("root", "123", "sshkey_file", "scriptDir", "bootdir",
                               "setupAgentFile", "ambariServer", "centos6",
                               None, "8440", "root")
    bootstrap_obj = Bootstrap("hostname", shared_state)
    getOsCheckScript_mock.return_value = "OsCheckScript"
    getOsCheckScriptRemoteLocation_mock.return_value = "OsCheckScriptRemoteLocation"
    expected = 42
    init_mock.return_value = None
    run_mock.return_value = expected
    res = bootstrap_obj.copyOsCheckScript()
    self.assertEquals(res, expected)
    input_file = str(init_mock.call_args[0][4])
    remote_file = str(init_mock.call_args[0][5])
    self.assertEqual(input_file, "OsCheckScript")
    self.assertEqual(remote_file, "OsCheckScriptRemoteLocation")


  @patch.object(BootstrapDefault, "getRemoteName")
  @patch.object(BootstrapDefault, "hasPassword")
  @patch.object(OSCheck, "is_suse_family")
  @patch.object(OSCheck, "is_ubuntu_family")
  @patch.object(OSCheck, "is_redhat_family")
  def test_getRepoFile(self, is_redhat_family, is_ubuntu_family, is_suse_family, hasPassword_mock, getRemoteName_mock):
    shared_state = SharedState("root", "123", "sshkey_file", "scriptDir", "bootdir",
                               "setupAgentFile", "ambariServer", "centos6",
                               None, "8440", "root")
    is_redhat_family.return_value = True
    is_ubuntu_family.return_value = False
    is_suse_family.return_value = False
    bootstrap_obj = Bootstrap("hostname", shared_state)
    # Without password
    hasPassword_mock.return_value = False
    getRemoteName_mock.return_value = "RemoteName"
    rf = bootstrap_obj.getMoveRepoFileCommand("target")
    self.assertEquals(rf, "/var/lib/ambari-agent/tmp/ambari-sudo.sh mv RemoteName target/ambari.repo")
    # With password
    hasPassword_mock.return_value = True
    getRemoteName_mock.return_value = "RemoteName"
    rf = bootstrap_obj.getMoveRepoFileCommand("target")
    self.assertEquals(rf, "/var/lib/ambari-agent/tmp/ambari-sudo.sh -S mv RemoteName target/ambari.repo < RemoteName")

  @patch("os.path.exists")
  @patch.object(OSCheck, "is_suse_family")
  @patch.object(OSCheck, "is_ubuntu_family")
  @patch.object(OSCheck, "is_redhat_family")
  @patch.object(BootstrapDefault, "getMoveRepoFileCommand")
  @patch.object(BootstrapDefault, "getRepoDir")
  @patch.object(BootstrapDefault, "getRepoFile")