Merge pull request #61 from dclayton-godaddy/tty-input

switch to tty input

Makefile

@@ -14,8 +14,10 @@ flake8:
 coverage:
 	PYTHONPATH=$(PYTHONPATH) pipenv run py.test --cov-config .coveragerc --verbose --cov-report term --cov-report xml --cov-report html --cov=aws_okta_processor tests
 
-publish:
+package:
 	python setup.py sdist bdist_wheel
+
+publish: package
 	twine check dist/*
 	twine upload dist/*
 	rm -fr build dist .egg requests.egg-info

aws_okta_processor/core/fetcher.py

@@ -8,7 +8,7 @@
 from hashlib import sha1
 from aws_okta_processor.core.okta import Okta
 from botocore.credentials import CachedCredentialFetcher
-from aws_okta_processor.core.print_tty import print_tty
+from aws_okta_processor.core.tty import print_tty
 
 
 class SAMLFetcher(CachedCredentialFetcher):

aws_okta_processor/core/okta.py

@@ -14,7 +14,7 @@
 from requests import ConnectTimeout
 from requests import ConnectionError
 from collections import OrderedDict
-from aws_okta_processor.core.print_tty import print_tty
+from aws_okta_processor.core.tty import print_tty, input_tty
 from six import add_metaclass
 
 
@@ -73,23 +73,23 @@ def __init__(
 
         if not self.organization:
             print_tty(string="Organization: ", newline=False)
-            self.organization = input()
+            self.organization = input_tty()
 
         if not self.user_name:
             print_tty(string="UserName: ", newline=False)
-            self.user_name = input()
+            self.user_name = input_tty()
 
         if not self.okta_session_id:
             if not self.user_name:
                 print_tty(string="UserName: ", newline=False)
-                self.user_name = input()
+                self.user_name = input_tty()
 
             if not user_pass:
                 user_pass = getpass.getpass()
 
             if not self.organization:
                 print_tty(string="Organization: ", newline=False)
-                self.organization = input()
+                self.organization = input_tty()
 
             self.okta_single_use_token = self.get_okta_single_use_token(
                 user_name=self.user_name,
@@ -450,7 +450,7 @@ def __init__(self, link=None):
     @staticmethod
     def payload():
         print_tty("Token: ", newline=False)
-        return {"passCode": input()}
+        return {"passCode": input_tty()}
 
     def retry(self, response):
         return False
@@ -465,7 +465,7 @@ def __init__(self, link=None):
     @staticmethod
     def payload():
         print_tty("Hardware Token: ", newline=False)
-        return {"passCode": input()}
+        return {"passCode": input_tty()}
 
     def retry(self, response):
         return False

aws_okta_processor/core/prompt.py

@@ -2,7 +2,7 @@
 import six
 
 from collections.abc import Mapping
-from aws_okta_processor.core.print_tty import print_tty
+from aws_okta_processor.core.tty import print_tty, input_tty
 
 
 BAD_INPUT_MESSAGE = "WARNING: Please supply a value from 1 to {}!"
@@ -66,7 +66,7 @@ def get_selection(options=None):
     print_tty("Selection: ", newline=False)
 
     try:
-        selection = input()
+        selection = input_tty()
     except KeyboardInterrupt:
         print_tty()
         sys.exit(1)

aws_okta_processor/core/saml.py

@@ -8,7 +8,7 @@
 
 from bs4 import BeautifulSoup
 from collections import OrderedDict
-from aws_okta_processor.core.print_tty import print_tty
+from aws_okta_processor.core.tty import print_tty
 
 
 SAML_ATTRIBUTE = '{urn:oasis:names:tc:SAML:2.0:assertion}Attribute'

aws_okta_processor/core/print_tty.py → aws_okta_processor/core/tty.py

@@ -8,6 +8,54 @@
 import sys
 
 
+def import_msvcrt():
+    import msvcrt
+    return msvcrt
+
+
+def input_tty():
+    try:
+        msvcrt = import_msvcrt()
+    except ImportError:
+        return unix_input_tty()
+    else:
+        return win_input_tty(msvcrt)
+
+
+def unix_input_tty():
+    with contextlib2.ExitStack() as stack:
+        try:
+            fd = os.open('/dev/tty', os.O_RDWR | os.O_NOCTTY)
+            tty = io.FileIO(fd, 'r+')
+            stack.enter_context(tty)
+            input = io.TextIOWrapper(tty)
+            stack.enter_context(input)
+        except OSError:
+            stack.close()
+            input = sys.stdin
+
+        line = input.readline()
+        if line[-1] == '\n':
+            line = line[:-1]
+        return line
+
+
+def win_input_tty(msvcrt):
+    pw = ""
+    while 1:
+        c = msvcrt.getwch()
+        if c == '\r' or c == '\n':
+            break
+        if c == '\003':
+            raise KeyboardInterrupt
+        if c == '\b':
+            pw = pw[:-1]
+        else:
+            pw = pw + c
+
+    return pw
+
+
 def unix_print_tty(string='', indents=0, newline=True):
     with contextlib2.ExitStack() as stack:
         string = indent(indents) + string
@@ -65,11 +113,6 @@ def indent(indents=None):
     return indent
 
 
-def import_msvcrt():
-    import msvcrt
-    return msvcrt
-
-
 def print_tty(string='', indents=0, newline=True, silent=False):
     try:
         msvcrt = import_msvcrt()

tests/core/test_fetcher.py

@@ -49,7 +49,7 @@ def test_get_app_roles(self, mock_get_app_roles):
     @patch("boto3.client")
     @patch('aws_okta_processor.core.fetcher.print_tty')
     @patch('aws_okta_processor.core.fetcher.prompt.print_tty')
-    @patch('aws_okta_processor.core.fetcher.prompt.input', return_value='1')
+    @patch('aws_okta_processor.core.fetcher.prompt.input_tty', return_value='1')
     @patch('aws_okta_processor.core.fetcher.Okta')
     def test_fetcher_should_filter_accounts(
             self,
@@ -104,7 +104,7 @@ def assume_role_side_effect(*args, **kwargs):
     @patch("boto3.client")
     @patch('aws_okta_processor.core.fetcher.print_tty')
     @patch('aws_okta_processor.core.fetcher.prompt.print_tty')
-    @patch('aws_okta_processor.core.fetcher.prompt.input', return_value='1')
+    @patch('aws_okta_processor.core.fetcher.prompt.input_tty', return_value='1')
     @patch('aws_okta_processor.core.fetcher.Okta')
     def test_fetcher_should_prompt_all_accounts(
             self,
@@ -155,4 +155,4 @@ def assume_role_side_effect(*args, **kwargs):
             call('Account: 2', indents=0),
             call('[ 3 ] Role-One', indents=1),
             call('Selection: ', newline=False)
-        ])
+        ])

tests/core/test_okta.py

@@ -252,7 +252,7 @@ def test_okta_mfa_push_challenge(
         self.assertEqual(okta.organization, "organization.okta.com")
         self.assertEqual(okta.okta_session_id, "session_token")
 
-    @patch('aws_okta_processor.core.okta.input')
+    @patch('aws_okta_processor.core.okta.input_tty')
     @patch('aws_okta_processor.core.okta.os.chmod')
     @patch('aws_okta_processor.core.okta.open')
     @patch('aws_okta_processor.core.okta.os.makedirs')
@@ -298,7 +298,7 @@ def test_okta_mfa_totp_challenge(
 
     @patch('aws_okta_processor.core.okta.Okta.get_okta_single_use_token')
     @patch('aws_okta_processor.core.okta.Okta.create_and_store_okta_session')
-    @patch('aws_okta_processor.core.okta.input')
+    @patch('aws_okta_processor.core.okta.input_tty')
     def test_read_aop_from_okta_session_should_read_aop_options(
         self,
         mock_input,
@@ -326,7 +326,7 @@ def test_read_aop_from_okta_session_should_read_aop_options(
     @patch('aws_okta_processor.core.okta.Okta.get_cache_file_path', return_value='/tmp/test.json')
     @patch('aws_okta_processor.core.okta.Okta.get_okta_single_use_token')
     @patch('aws_okta_processor.core.okta.Okta.create_and_store_okta_session')
-    @patch('aws_okta_processor.core.okta.input')
+    @patch('aws_okta_processor.core.okta.input_tty')
     @patch('builtins.open', new_callable=mock_open)
     def test_set_okta_session_should_write_session_data(
         self,
@@ -368,7 +368,7 @@ def test_set_okta_session_should_write_session_data(
             call('}')
         ])
 
-    @patch('aws_okta_processor.core.okta.input')
+    @patch('aws_okta_processor.core.okta.input_tty')
     @patch('aws_okta_processor.core.okta.os.chmod')
     @patch('aws_okta_processor.core.okta.open')
     @patch('aws_okta_processor.core.okta.os.makedirs')
@@ -412,7 +412,7 @@ def test_okta_mfa_hardware_token_challenge(
         self.assertEqual(okta.organization, "organization.okta.com")
         self.assertEqual(okta.okta_session_id, "session_token")
 
-    @patch('aws_okta_processor.core.prompt.input')
+    @patch('aws_okta_processor.core.prompt.input_tty')
     @patch('aws_okta_processor.core.okta.os.chmod')
     @patch('aws_okta_processor.core.okta.open')
     @patch('aws_okta_processor.core.okta.os.makedirs')

tests/core/test_prompt.py

@@ -65,7 +65,7 @@ def test_get_item_config_no_match(self, mock_sys, mock_print_tty):
         )
 
     @patch('aws_okta_processor.core.prompt.print_tty')
-    @patch('aws_okta_processor.core.prompt.input')
+    @patch('aws_okta_processor.core.prompt.input_tty')
     def test_get_selection(self, mock_input, mock_print_tty):
         mock_input.return_value = 1
         options = ["one", "two"]
@@ -77,7 +77,7 @@ def test_get_selection(self, mock_input, mock_print_tty):
         )
 
     @patch('aws_okta_processor.core.prompt.print_tty')
-    @patch('aws_okta_processor.core.prompt.input')
+    @patch('aws_okta_processor.core.prompt.input_tty')
     def test_get_selection_bad_input(self, mock_input, mock_print_tty):
         mock_input.side_effect = ["bad_input", 2]
         options = ["one", "two"]
@@ -89,7 +89,7 @@ def test_get_selection_bad_input(self, mock_input, mock_print_tty):
         )
 
     @patch('aws_okta_processor.core.prompt.print_tty')
-    @patch('aws_okta_processor.core.prompt.input')
+    @patch('aws_okta_processor.core.prompt.input_tty')
     def test_get_selection_bad_int(self, mock_input, mock_print_tty):
         mock_input.side_effect = [0, 2]
         options = ["one", "two"]
@@ -101,7 +101,7 @@ def test_get_selection_bad_int(self, mock_input, mock_print_tty):
         )
 
     @patch('aws_okta_processor.core.prompt.print_tty')
-    @patch('aws_okta_processor.core.prompt.input')
+    @patch('aws_okta_processor.core.prompt.input_tty')
     def test_get_selection_no_input(self, mock_input, mock_print_tty):
         mock_input.side_effect = [SyntaxError, 2]
         options = ["one", "two"]
@@ -114,7 +114,7 @@ def test_get_selection_no_input(self, mock_input, mock_print_tty):
 
     @patch('aws_okta_processor.core.prompt.print_tty')
     @patch('aws_okta_processor.core.prompt.sys')
-    @patch('aws_okta_processor.core.prompt.input')
+    @patch('aws_okta_processor.core.prompt.input_tty')
     def test_get_selection_interrupt(self, mock_input, mock_sys, mock_print_tty): # noqa
         mock_input.side_effect = [KeyboardInterrupt, 2]
         options = ["one"]