# -*- coding: utf-8 -*-
# Copyright 2021 Google Inc. All Rights Reserved.
#
# 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.
"""Tests for stet_util.py."""

from __future__ import absolute_import
from __future__ import print_function
from __future__ import division
from __future__ import unicode_literals

import os
import shutil

from gslib import storage_url
from gslib.tests import testcase
from gslib.tests import util
from gslib.tests.util import unittest
from gslib.utils import execution_util
from gslib.utils import stet_util

import mock


class TestStetUtil(testcase.GsUtilUnitTestCase):
  """Test STET utils."""

  @mock.patch.object(execution_util, 'ExecuteExternalCommand')
  def test_stet_upload_uses_binary_and_config_from_boto(
      self, mock_execute_external_command):
    fake_config_path = self.CreateTempFile()
    mock_execute_external_command.return_value = ('stdout', 'stderr')
    mock_logger = mock.Mock()

    source_url = storage_url.StorageUrlFromString('in')
    destination_url = storage_url.StorageUrlFromString('gs://bucket/obj')
    with util.SetBotoConfigForTest([
        ('GSUtil', 'stet_binary_path', 'fake_binary_path'),
        ('GSUtil', 'stet_config_path', fake_config_path),
    ]):
      out_file_url = stet_util.encrypt_upload(source_url, destination_url,
                                              mock_logger)

    self.assertEqual(out_file_url,
                     storage_url.StorageUrlFromString('in_.stet_tmp'))
    mock_execute_external_command.assert_called_once_with([
        'fake_binary_path',
        'encrypt',
        '--config-file={}'.format(fake_config_path),
        '--blob-id=gs://bucket/obj',
        'in',
        'in_.stet_tmp',
    ])
    mock_logger.debug.assert_called_once_with('stderr')

  @mock.patch.object(execution_util, 'ExecuteExternalCommand')
  def test_stet_upload_runs_with_binary_from_path_with_correct_settings(
      self, mock_execute_external_command):
    fake_config_path = self.CreateTempFile()
    temporary_path_directory = self.CreateTempDir()
    fake_stet_binary_path = self.CreateTempFile(tmpdir=temporary_path_directory,
                                                file_name='stet')
    previous_path = os.getenv('PATH')
    os.environ['PATH'] += os.path.pathsep + temporary_path_directory

    mock_execute_external_command.return_value = ('stdout', 'stderr')
    mock_logger = mock.Mock()

    source_url = storage_url.StorageUrlFromString('in')
    destination_url = storage_url.StorageUrlFromString('gs://bucket/obj')
    with util.SetBotoConfigForTest([
        ('GSUtil', 'stet_binary_path', None),
        ('GSUtil', 'stet_config_path', fake_config_path),
    ]):
      out_file_url = stet_util.encrypt_upload(source_url, destination_url,
                                              mock_logger)

    self.assertEqual(out_file_url,
                     storage_url.StorageUrlFromString('in_.stet_tmp'))
    mock_execute_external_command.assert_called_once_with([
        fake_stet_binary_path,
        'encrypt',
        '--config-file={}'.format(fake_config_path),
        '--blob-id=gs://bucket/obj',
        'in',
        'in_.stet_tmp',
    ])
    mock_logger.debug.assert_called_once_with('stderr')

    os.environ['PATH'] = previous_path

  @mock.patch.object(execution_util, 'ExecuteExternalCommand')
  def test_stet_upload_uses_config_from_default_path_with_correct_settings(
      self, mock_execute_external_command):
    mock_execute_external_command.return_value = ('stdout', 'stderr')
    mock_logger = mock.Mock()

    source_url = storage_url.StorageUrlFromString('in')
    destination_url = storage_url.StorageUrlFromString('gs://bucket/obj')
    with util.SetBotoConfigForTest([
        ('GSUtil', 'stet_binary_path', 'fake_binary_path'),
        ('GSUtil', 'stet_config_path', None),
    ]):
      with mock.patch.object(os.path,
                             'exists',
                             new=mock.Mock(return_value=True)):
        out_file_url = stet_util.encrypt_upload(source_url, destination_url,
                                                mock_logger)

    self.assertEqual(out_file_url,
                     storage_url.StorageUrlFromString('in_.stet_tmp'))
    mock_execute_external_command.assert_called_once_with([
        'fake_binary_path',
        'encrypt',
        '--config-file={}'.format(
            os.path.expanduser(stet_util.DEFAULT_STET_CONFIG_PATH)),
        '--blob-id=gs://bucket/obj',
        'in',
        'in_.stet_tmp',
    ])
    mock_logger.debug.assert_called_once_with('stderr')

  @mock.patch.object(shutil, 'move')
  @mock.patch.object(execution_util, 'ExecuteExternalCommand')
  def test_stet_download_runs_binary_and_replaces_temp_file(
      self, mock_execute_external_command, mock_move):
    fake_config_path = self.CreateTempFile()
    mock_execute_external_command.return_value = ('stdout', 'stderr')
    mock_logger = mock.Mock()

    source_url = storage_url.StorageUrlFromString('gs://bucket/obj')
    destination_url = storage_url.StorageUrlFromString('out')
    with util.SetBotoConfigForTest([
        ('GSUtil', 'stet_binary_path', 'fake_binary_path'),
        ('GSUtil', 'stet_config_path', fake_config_path),
    ]):
      stet_util.decrypt_download(source_url, destination_url, mock_logger)

    mock_execute_external_command.assert_called_once_with([
        'fake_binary_path', 'decrypt',
        '--config-file={}'.format(fake_config_path),
        '--blob-id=gs://bucket/obj', 'out', 'out_.stet_tmp'
    ])
    mock_logger.debug.assert_called_once_with('stderr')
    mock_move.assert_called_once_with('out_.stet_tmp', 'out')

  @mock.patch.object(stet_util,
                     '_get_stet_binary_from_path',
                     new=mock.Mock(return_value=None))
  def test_stet_util_errors_if_no_binary(self):
    fake_config_path = self.CreateTempFile()
    source_url = storage_url.StorageUrlFromString('in')
    destination_url = storage_url.StorageUrlFromString('gs://bucket/obj')
    with util.SetBotoConfigForTest([
        ('GSUtil', 'stet_binary_path', None),
        ('GSUtil', 'stet_config_path', fake_config_path),
    ]):
      with self.assertRaises(KeyError):
        stet_util.encrypt_upload(source_url, destination_url, None)

  def test_stet_util_errors_if_no_config(self):
    source_url = storage_url.StorageUrlFromString('in')
    destination_url = storage_url.StorageUrlFromString('gs://bucket/obj')
    with util.SetBotoConfigForTest([
        ('GSUtil', 'stet_binary_path', 'fake_binary_path'),
        ('GSUtil', 'stet_config_path', None),
    ]):
      with mock.patch.object(os.path,
                             'exists',
                             new=mock.Mock(return_value=False)):
        with self.assertRaises(KeyError):
          stet_util.encrypt_upload(source_url, destination_url, None)

  @mock.patch.object(os.path, 'expanduser', autospec=True)
  @mock.patch.object(execution_util,
                     'ExecuteExternalCommand',
                     new=mock.Mock(return_value=('stdout', 'stderr')))
  def test_stet_util_expands_home_directory_symbol(self, mock_expanduser):
    fake_config_path = self.CreateTempFile()
    source_url = storage_url.StorageUrlFromString('in')
    destination_url = storage_url.StorageUrlFromString('gs://bucket/obj')
    with util.SetBotoConfigForTest([
        ('GSUtil', 'stet_binary_path', 'fake_binary_path'),
        ('GSUtil', 'stet_config_path', fake_config_path),
    ]):
      stet_util.encrypt_upload(source_url, destination_url, mock.Mock())
    mock_expanduser.assert_has_calls(
        [mock.call('fake_binary_path'),
         mock.call(fake_config_path)])

[((38, 3, 38, 62), 'mock.patch.object', 'mock.patch.object', ({(38, 21, 38, 35): 'execution_util', (38, 37, 38, 61): '"""ExecuteExternalCommand"""'}, {}), "(execution_util, 'ExecuteExternalCommand')", False, 'import mock\n'), ((66, 3, 66, 62), 'mock.patch.object', 'mock.patch.object', ({(66, 21, 66, 35): 'execution_util', (66, 37, 66, 61): '"""ExecuteExternalCommand"""'}, {}), "(execution_util, 'ExecuteExternalCommand')", False, 'import mock\n'), ((102, 3, 102, 62), 'mock.patch.object', 'mock.patch.object', ({(102, 21, 102, 35): 'execution_util', (102, 37, 102, 61): '"""ExecuteExternalCommand"""'}, {}), "(execution_util, 'ExecuteExternalCommand')", False, 'import mock\n'), ((133, 3, 133, 36), 'mock.patch.object', 'mock.patch.object', ({(133, 21, 133, 27): 'shutil', (133, 29, 133, 35): '"""move"""'}, {}), "(shutil, 'move')", False, 'import mock\n'), ((134, 3, 134, 62), 'mock.patch.object', 'mock.patch.object', ({(134, 21, 134, 35): 'execution_util', (134, 37, 134, 61): '"""ExecuteExternalCommand"""'}, {}), "(execution_util, 'ExecuteExternalCommand')", False, 'import mock\n'), ((184, 3, 184, 58), 'mock.patch.object', 'mock.patch.object', (), '', False, 'import mock\n'), ((43, 18, 43, 29), 'mock.Mock', 'mock.Mock', ({}, {}), '()', False, 'import mock\n'), ((45, 17, 45, 55), 'gslib.storage_url.StorageUrlFromString', 'storage_url.StorageUrlFromString', ({(45, 50, 45, 54): '"""in"""'}, {}), "('in')", False, 'from gslib import storage_url\n'), ((46, 22, 46, 73), 'gslib.storage_url.StorageUrlFromString', 'storage_url.StorageUrlFromString', ({(46, 55, 46, 72): '"""gs://bucket/obj"""'}, {}), "('gs://bucket/obj')", False, 'from gslib import storage_url\n'), ((73, 20, 73, 37), 'os.getenv', 'os.getenv', ({(73, 30, 73, 36): '"""PATH"""'}, {}), "('PATH')", False, 'import os\n'), ((77, 18, 77, 29), 'mock.Mock', 'mock.Mock', ({}, {}), '()', False, 'import mock\n'), ((79, 17, 79, 55), 'gslib.storage_url.StorageUrlFromString', 'storage_url.StorageUrlFromString', ({(79, 50, 79, 54): '"""in"""'}, {}), "('in')", False, 'from gslib import storage_url\n'), ((80, 22, 80, 73), 'gslib.storage_url.StorageUrlFromString', 'storage_url.StorageUrlFromString', ({(80, 55, 80, 72): '"""gs://bucket/obj"""'}, {}), "('gs://bucket/obj')", False, 'from gslib import storage_url\n'), ((106, 18, 106, 29), 'mock.Mock', 'mock.Mock', ({}, {}), '()', False, 'import mock\n'), ((108, 17, 108, 55), 'gslib.storage_url.StorageUrlFromString', 'storage_url.StorageUrlFromString', ({(108, 50, 108, 54): '"""in"""'}, {}), "('in')", False, 'from gslib import storage_url\n'), ((109, 22, 109, 73), 'gslib.storage_url.StorageUrlFromString', 'storage_url.StorageUrlFromString', ({(109, 55, 109, 72): '"""gs://bucket/obj"""'}, {}), "('gs://bucket/obj')", False, 'from gslib import storage_url\n'), ((139, 18, 139, 29), 'mock.Mock', 'mock.Mock', ({}, {}), '()', False, 'import mock\n'), ((141, 17, 141, 68), 'gslib.storage_url.StorageUrlFromString', 'storage_url.StorageUrlFromString', ({(141, 50, 141, 67): '"""gs://bucket/obj"""'}, {}), "('gs://bucket/obj')", False, 'from gslib import storage_url\n'), ((142, 22, 142, 61), 'gslib.storage_url.StorageUrlFromString', 'storage_url.StorageUrlFromString', ({(142, 55, 142, 60): '"""out"""'}, {}), "('out')", False, 'from gslib import storage_url\n'), ((162, 17, 162, 55), 'gslib.storage_url.StorageUrlFromString', 'storage_url.StorageUrlFromString', ({(162, 50, 162, 54): '"""in"""'}, {}), "('in')", False, 'from gslib import storage_url\n'), ((163, 22, 163, 73), 'gslib.storage_url.StorageUrlFromString', 'storage_url.StorageUrlFromString', ({(163, 55, 163, 72): '"""gs://bucket/obj"""'}, {}), "('gs://bucket/obj')", False, 'from gslib import storage_url\n'), ((172, 17, 172, 55), 'gslib.storage_url.StorageUrlFromString', 'storage_url.StorageUrlFromString', ({(172, 50, 172, 54): '"""in"""'}, {}), "('in')", False, 'from gslib import storage_url\n'), ((173, 22, 173, 73), 'gslib.storage_url.StorageUrlFromString', 'storage_url.StorageUrlFromString', ({(173, 55, 173, 72): '"""gs://bucket/obj"""'}, {}), "('gs://bucket/obj')", False, 'from gslib import storage_url\n'), ((190, 17, 190, 55), 'gslib.storage_url.StorageUrlFromString', 'storage_url.StorageUrlFromString', ({(190, 50, 190, 54): '"""in"""'}, {}), "('in')", False, 'from gslib import storage_url\n'), ((191, 22, 191, 73), 'gslib.storage_url.StorageUrlFromString', 'storage_url.StorageUrlFromString', ({(191, 55, 191, 72): '"""gs://bucket/obj"""'}, {}), "('gs://bucket/obj')", False, 'from gslib import storage_url\n'), ((47, 9, 50, 6), 'gslib.tests.util.SetBotoConfigForTest', 'util.SetBotoConfigForTest', ({(47, 35, 50, 5): "[('GSUtil', 'stet_binary_path', 'fake_binary_path'), ('GSUtil',\n    'stet_config_path', fake_config_path)]"}, {}), "([('GSUtil', 'stet_binary_path',\n    'fake_binary_path'), ('GSUtil', 'stet_config_path', fake_config_path)])", False, 'from gslib.tests import util\n'), ((51, 21, 52, 58), 'gslib.utils.stet_util.encrypt_upload', 'stet_util.encrypt_upload', ({(51, 46, 51, 56): 'source_url', (51, 58, 51, 73): 'destination_url', (52, 46, 52, 57): 'mock_logger'}, {}), '(source_url, destination_url, mock_logger)', False, 'from gslib.utils import stet_util\n'), ((55, 21, 55, 69), 'gslib.storage_url.StorageUrlFromString', 'storage_url.StorageUrlFromString', ({(55, 54, 55, 68): '"""in_.stet_tmp"""'}, {}), "('in_.stet_tmp')", False, 'from gslib import storage_url\n'), ((81, 9, 84, 6), 'gslib.tests.util.SetBotoConfigForTest', 'util.SetBotoConfigForTest', ({(81, 35, 84, 5): "[('GSUtil', 'stet_binary_path', None), ('GSUtil', 'stet_config_path',\n    fake_config_path)]"}, {}), "([('GSUtil', 'stet_binary_path', None), ('GSUtil',\n    'stet_config_path', fake_config_path)])", False, 'from gslib.tests import util\n'), ((85, 21, 86, 58), 'gslib.utils.stet_util.encrypt_upload', 'stet_util.encrypt_upload', ({(85, 46, 85, 56): 'source_url', (85, 58, 85, 73): 'destination_url', (86, 46, 86, 57): 'mock_logger'}, {}), '(source_url, destination_url, mock_logger)', False, 'from gslib.utils import stet_util\n'), ((89, 21, 89, 69), 'gslib.storage_url.StorageUrlFromString', 'storage_url.StorageUrlFromString', ({(89, 54, 89, 68): '"""in_.stet_tmp"""'}, {}), "('in_.stet_tmp')", False, 'from gslib import storage_url\n'), ((110, 9, 113, 6), 'gslib.tests.util.SetBotoConfigForTest', 'util.SetBotoConfigForTest', ({(110, 35, 113, 5): "[('GSUtil', 'stet_binary_path', 'fake_binary_path'), ('GSUtil',\n    'stet_config_path', None)]"}, {}), "([('GSUtil', 'stet_binary_path',\n    'fake_binary_path'), ('GSUtil', 'stet_config_path', None)])", False, 'from gslib.tests import util\n'), ((121, 21, 121, 69), 'gslib.storage_url.StorageUrlFromString', 'storage_url.StorageUrlFromString', ({(121, 54, 121, 68): '"""in_.stet_tmp"""'}, {}), "('in_.stet_tmp')", False, 'from gslib import storage_url\n'), ((143, 9, 146, 6), 'gslib.tests.util.SetBotoConfigForTest', 'util.SetBotoConfigForTest', ({(143, 35, 146, 5): "[('GSUtil', 'stet_binary_path', 'fake_binary_path'), ('GSUtil',\n    'stet_config_path', fake_config_path)]"}, {}), "([('GSUtil', 'stet_binary_path',\n    'fake_binary_path'), ('GSUtil', 'stet_config_path', fake_config_path)])", False, 'from gslib.tests import util\n'), ((147, 6, 147, 74), 'gslib.utils.stet_util.decrypt_download', 'stet_util.decrypt_download', ({(147, 33, 147, 43): 'source_url', (147, 45, 147, 60): 'destination_url', (147, 62, 147, 73): 'mock_logger'}, {}), '(source_url, destination_url, mock_logger)', False, 'from gslib.utils import stet_util\n'), ((164, 9, 167, 6), 'gslib.tests.util.SetBotoConfigForTest', 'util.SetBotoConfigForTest', ({(164, 35, 167, 5): "[('GSUtil', 'stet_binary_path', None), ('GSUtil', 'stet_config_path',\n    fake_config_path)]"}, {}), "([('GSUtil', 'stet_binary_path', None), ('GSUtil',\n    'stet_config_path', fake_config_path)])", False, 'from gslib.tests import util\n'), ((159, 25, 159, 53), 'mock.Mock', 'mock.Mock', (), '', False, 'import mock\n'), ((174, 9, 177, 6), 'gslib.tests.util.SetBotoConfigForTest', 'util.SetBotoConfigForTest', ({(174, 35, 177, 5): "[('GSUtil', 'stet_binary_path', 'fake_binary_path'), ('GSUtil',\n    'stet_config_path', None)]"}, {}), "([('GSUtil', 'stet_binary_path',\n    'fake_binary_path'), ('GSUtil', 'stet_config_path', None)])", False, 'from gslib.tests import util\n'), ((192, 9, 195, 6), 'gslib.tests.util.SetBotoConfigForTest', 'util.SetBotoConfigForTest', ({(192, 35, 195, 5): "[('GSUtil', 'stet_binary_path', 'fake_binary_path'), ('GSUtil',\n    'stet_config_path', fake_config_path)]"}, {}), "([('GSUtil', 'stet_binary_path',\n    'fake_binary_path'), ('GSUtil', 'stet_config_path', fake_config_path)])", False, 'from gslib.tests import util\n'), ((187, 25, 187, 69), 'mock.Mock', 'mock.Mock', (), '', False, 'import mock\n'), ((117, 23, 118, 60), 'gslib.utils.stet_util.encrypt_upload', 'stet_util.encrypt_upload', ({(117, 48, 117, 58): 'source_url', (117, 60, 117, 75): 'destination_url', (118, 48, 118, 59): 'mock_logger'}, {}), '(source_url, destination_url, mock_logger)', False, 'from gslib.utils import stet_util\n'), ((169, 8, 169, 67), 'gslib.utils.stet_util.encrypt_upload', 'stet_util.encrypt_upload', ({(169, 33, 169, 43): 'source_url', (169, 45, 169, 60): 'destination_url', (169, 62, 169, 66): 'None'}, {}), '(source_url, destination_url, None)', False, 'from gslib.utils import stet_util\n'), ((196, 60, 196, 71), 'mock.Mock', 'mock.Mock', ({}, {}), '()', False, 'import mock\n'), ((198, 9, 198, 38), 'mock.call', 'mock.call', ({(198, 19, 198, 37): '"""fake_binary_path"""'}, {}), "('fake_binary_path')", False, 'import mock\n'), ((199, 9, 199, 36), 'mock.call', 'mock.call', ({(199, 19, 199, 35): 'fake_config_path'}, {}), '(fake_config_path)', False, 'import mock\n'), ((126, 12, 126, 66), 'os.path.expanduser', 'os.path.expanduser', ({(126, 31, 126, 65): 'stet_util.DEFAULT_STET_CONFIG_PATH'}, {}), '(stet_util.DEFAULT_STET_CONFIG_PATH)', False, 'import os\n'), ((182, 10, 182, 69), 'gslib.utils.stet_util.encrypt_upload', 'stet_util.encrypt_upload', ({(182, 35, 182, 45): 'source_url', (182, 47, 182, 62): 'destination_url', (182, 64, 182, 68): 'None'}, {}), '(source_url, destination_url, None)', False, 'from gslib.utils import stet_util\n'), ((116, 33, 116, 61), 'mock.Mock', 'mock.Mock', (), '', False, 'import mock\n'), ((180, 33, 180, 62), 'mock.Mock', 'mock.Mock', (), '', False, 'import mock\n')]