Always use own version of session-manager-plugin (#42)

xen0l · web-flow · commit 56f3db6b700f · 2019-08-20T02:37:26.000+02:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,6 +2,7 @@
 ------------------
 
 * Use default region and profile in list command
+* Always use own version of session-manager-plugin if available
 
 0.4.1 (2019-08-12)
 ------------------
diff --git a/aws_gate/bootstrap.py b/aws_gate/bootstrap.py
@@ -6,14 +6,11 @@
 import platform
 import requests
 
-from aws_gate.config import DEFAULT_GATE_BIN_PATH, PLUGIN_INSTALL_PATH, PLUGIN_NAME
+from aws_gate.constants import MAC_PLUGIN_URL, DEFAULT_GATE_BIN_PATH, PLUGIN_INSTALL_PATH, PLUGIN_NAME
 from aws_gate.exceptions import UnsupportedPlatormError
 from aws_gate.utils import execute
 
 
-MAC_PLUGIN_URL = 'https://s3.amazonaws.com/session-manager-downloads/plugin/latest/mac/sessionmanager-bundle.zip'
-
-
 logger = logging.getLogger(__name__)
 
 
diff --git a/aws_gate/config.py b/aws_gate/config.py
@@ -6,15 +6,9 @@
 from yaml.parser import ParserError
 from marshmallow import Schema, fields, post_load, ValidationError, validates_schema
 
+from aws_gate.constants import DEFAULT_GATE_CONFIG_PATH, DEFAULT_GATE_CONFIGD_PATH
 from aws_gate.utils import is_existing_profile, is_existing_region
 
-DEFAULT_GATE_DIR = os.path.expanduser('~/.aws-gate')
-DEFAULT_GATE_CONFIG_PATH = os.path.join(DEFAULT_GATE_DIR, 'config')
-DEFAULT_GATE_CONFIGD_PATH = os.path.join(DEFAULT_GATE_DIR, 'config.d')
-
-PLUGIN_NAME = 'session-manager-plugin'
-DEFAULT_GATE_BIN_PATH = os.path.join(DEFAULT_GATE_DIR, 'bin')
-PLUGIN_INSTALL_PATH = os.path.join(DEFAULT_GATE_BIN_PATH, PLUGIN_NAME)
 
 logger = logging.getLogger(__name__)
 
diff --git a/aws_gate/constants.py b/aws_gate/constants.py
@@ -0,0 +1,11 @@
+import os
+
+DEFAULT_GATE_DIR = os.path.expanduser('~/.aws-gate')
+DEFAULT_GATE_CONFIG_PATH = os.path.join(DEFAULT_GATE_DIR, 'config')
+DEFAULT_GATE_CONFIGD_PATH = os.path.join(DEFAULT_GATE_DIR, 'config.d')
+
+PLUGIN_NAME = 'session-manager-plugin'
+DEFAULT_GATE_BIN_PATH = os.path.join(DEFAULT_GATE_DIR, 'bin')
+PLUGIN_INSTALL_PATH = os.path.join(DEFAULT_GATE_BIN_PATH, PLUGIN_NAME)
+
+MAC_PLUGIN_URL = 'https://s3.amazonaws.com/session-manager-downloads/plugin/latest/mac/sessionmanager-bundle.zip'
diff --git a/aws_gate/decorators.py b/aws_gate/decorators.py
@@ -1,8 +1,8 @@
 import os
 import functools
 
+from aws_gate.constants import PLUGIN_INSTALL_PATH
 from aws_gate.utils import execute
-from aws_gate.config import PLUGIN_INSTALL_PATH
 
 
 def _plugin_exists(plugin_path):
diff --git a/aws_gate/session.py b/aws_gate/session.py
@@ -1,11 +1,11 @@
 import errno
 import json
 import logging
-import subprocess
 
 from aws_gate.query import query_instance
 from aws_gate.decorators import plugin_version, plugin_required
-from aws_gate.utils import get_aws_client, get_aws_resource, deferred_signals, is_existing_profile, is_existing_region
+from aws_gate.utils import get_aws_client, get_aws_resource, deferred_signals, is_existing_profile, \
+    is_existing_region, execute
 
 logger = logging.getLogger(__name__)
 
@@ -44,12 +44,9 @@ def terminate(self):
     def open(self):
         with deferred_signals():
             try:
-                logger.debug('Executing session-manager-plugin')
-                subprocess.check_call(['session-manager-plugin',
-                                       json.dumps(self._response),
-                                       self._region_name,
-                                       'StartSession'])
-                logger.debug('session-manager plugin finished successfully')
+                execute('session-manager-plugin', [json.dumps(self._response),
+                                                   self._region_name,
+                                                   'StartSession'])
             except OSError as ex:
                 if ex.errno == errno.ENOENT:
                     raise ValueError('session-manager-plugin cannot be found')
diff --git a/aws_gate/utils.py b/aws_gate/utils.py
@@ -6,6 +6,7 @@
 
 import boto3
 
+from aws_gate.constants import DEFAULT_GATE_BIN_PATH
 
 logger = logging.getLogger(__name__)
 
@@ -87,15 +88,18 @@ def deferred_signals(signal_list=None):
             signal.signal(deferred_signal, signal.SIG_DFL)
 
 
-def execute(path, args, capture_output=True):
-    ret = None
+def execute(cmd, args, **kwargs):
+    ret, result = None, None
+
+    env = DEFAULT_GATE_BIN_PATH + os.pathsep + os.environ['PATH']
+
     try:
-        logger.debug('Executing "%s"', ' '.join([path] + args))
-        result = subprocess.run([path] + args, capture_output=capture_output)
+        logger.debug('Executing "%s"', ' '.join([cmd] + args))
+        result = subprocess.run([cmd] + args, env={'PATH': env}, check=True, **kwargs)
     except subprocess.CalledProcessError as e:
-        logger.error('Command "%s" exited with %s', ' '.join([path] + args), e.returncode)
+        logger.error('Command "%s" exited with %s', ' '.join([cmd] + args), e.returncode)
 
-    if result.stdout:
+    if result and result.stdout:
         ret = result.stdout.decode()
         ret = ret.rstrip()
 
diff --git a/requirements/requirements_dev.txt b/requirements/requirements_dev.txt
@@ -6,3 +6,4 @@ pytest-flake8==1.0.4
 pytest-pylint==0.14.1
 pytest==5.0.1
 zest.releaser==6.19.0
+hypothesis==4.32.3
diff --git a/test/unit/test_bootstrap.py b/test/unit/test_bootstrap.py
@@ -1,8 +1,8 @@
 import unittest
 from unittest.mock import patch, MagicMock, call, mock_open
 
+from aws_gate.constants import DEFAULT_GATE_BIN_PATH, PLUGIN_INSTALL_PATH
 from aws_gate.exceptions import UnsupportedPlatormError
-from aws_gate.config import DEFAULT_GATE_BIN_PATH, PLUGIN_INSTALL_PATH
 from aws_gate.bootstrap import Plugin, MacPlugin, bootstrap
 
 
diff --git a/test/unit/test_config.py b/test/unit/test_config.py
@@ -4,8 +4,8 @@
 from unittest.mock import patch
 from marshmallow import ValidationError
 
-from aws_gate.config import GateConfig, EmptyConfigurationError, load_config_from_files, _locate_config_files, \
-    DEFAULT_GATE_CONFIG_PATH, DEFAULT_GATE_CONFIGD_PATH
+from aws_gate.constants import DEFAULT_GATE_CONFIG_PATH, DEFAULT_GATE_CONFIGD_PATH
+from aws_gate.config import GateConfig, EmptyConfigurationError, load_config_from_files, _locate_config_files
 
 
 class TestConfig(unittest.TestCase):
diff --git a/test/unit/test_session.py b/test/unit/test_session.py
@@ -26,9 +26,6 @@ def setUp(self):
             'TokenValue': 'randomtokenvalue'
         }
 
-#        patch('aws_gate.decorators._plugin_exists', return_value=True).start()
-#        patch('aws_gate.decorators.execute', return_value='1.1.23.0').start()
-
     def test_create_session(self):
         with patch.object(self.ssm, 'start_session', return_value=self.response):
             sess = Session(instance_id=self.instance_id, ssm=self.ssm)
@@ -46,14 +43,16 @@ def test_terminate_session(self):
             self.assertTrue(self.ssm.terminate_session.called)
 
     def test_open_session(self):
-        with patch('aws_gate.session.subprocess.check_call', return_value=True) as m:
+        mock_output = MagicMock(stdout=b'output')
+
+        with patch('aws_gate.session.execute', return_value=mock_output) as m:
             sess = Session(instance_id=self.instance_id, ssm=self.ssm)
             sess.open()
 
             self.assertTrue(m.called)
 
     def test_open_session_exception(self):
-        with patch('aws_gate.session.subprocess.check_call',
+        with patch('aws_gate.session.execute',
                    side_effect=OSError(errno.ENOENT, os.strerror(errno.ENOENT))):
             with self.assertRaises(ValueError):
                 sess = Session(instance_id=self.instance_id, ssm=self.ssm)
diff --git a/test/unit/test_utils.py b/test/unit/test_utils.py
@@ -1,8 +1,12 @@
+import subprocess
 import unittest
 from unittest.mock import patch, MagicMock, call
 
+from hypothesis import given
+from hypothesis.strategies import lists, text
+
 from aws_gate.utils import is_existing_profile, _create_aws_session, get_aws_client, get_aws_resource, \
-    AWS_REGIONS, is_existing_region
+    AWS_REGIONS, is_existing_region, execute
 
 
 # pylint: disable=too-few-public-methods
@@ -51,3 +55,17 @@ def test_get_aws_resource(self):
     def test_region_validation(self):
         self.assertTrue(is_existing_region(region_name=AWS_REGIONS[0]))
         self.assertFalse(is_existing_region(region_name='unknown-region-1'))
+
+    @given(text(), lists(text()))
+    def test_execute(self, cmd, args):
+        mock_output = MagicMock(stdout=b'output')
+
+        with patch('aws_gate.utils.subprocess.run', return_value=mock_output):
+            self.assertEqual(execute(cmd, args), 'output')
+
+    def test_execute_exception(self):
+        with patch('aws_gate.utils.subprocess.run',
+                   side_effect=subprocess.CalledProcessError(returncode=1, cmd='error')) as mock:
+            execute('/usr/bin/ls', ['-l'])
+
+            self.assertTrue(mock.called)
