Repository: byu-oit/awslogin

Issue:
record logged in account/role in a file
I would like to be able to add my last used account/role into my shell prompt for convenience. Something like:

PowerUser@prd-sim-appdev-oit-byu /opt/webroot$

I would like to be able to look somewhere or make a call to get the account and role information.

Discussion:

Patch:
diff --git a/README.rst b/README.rst
index b1a031d..8e13308 100644
--- a/README.rst
+++ b/README.rst
@@ -54,8 +54,15 @@ To use it:
    arguments as well.
 -  Run ``awslogin --profile <profile name>`` to specifiy an alternative
    profile
+-  Run ``awslogin --status`` for the current status of the default profile
+-  Run ``awslogin --status --profile=dev`` for the current status of the
+   dev profile
+-  Run ``awslogin --status --profile=all`` for the current status of the
+   all profiles
+-  Run ``awslogin --version`` to display the running version of awslogin
 -  Run ``awslogin -- --help`` for full help message
 
+
 Reporting bugs or requesting features
 -------------------------------------
 
diff --git a/VERSION b/VERSION
deleted file mode 100644
index af88ba8..0000000
--- a/VERSION
+++ /dev/null
@@ -1,1 +0,0 @@
-0.11.1
diff --git a/byu_awslogin/index.py b/byu_awslogin/index.py
index 16a27d0..1ddcdcb 100755
--- a/byu_awslogin/index.py
+++ b/byu_awslogin/index.py
@@ -4,6 +4,7 @@ import configparser
 import getpass
 import os
 import sys
+import datetime
 from os.path import expanduser
 
 import fire
@@ -23,8 +24,16 @@ from .adfs_auth import authenticate
 from .assume_role import ask_which_role_to_assume, assume_role
 from .roles import action_url_on_validation_success, retrieve_roles_page
 
+__VERSION__ = '0.11.2'
 
-def cli(account=None, role=None, profile='default'):
+
+def cli(account=None, role=None, profile='default', version=False, status=False):
+    if version:
+        print(f'awslogin version: {__VERSION__}')
+        return
+    if status:
+        get_status(aws_file('config'), profile)
+        return
     # Get the federated credentials from the user
     cached_netid = load_last_netid(aws_file('config'), profile)
     if cached_netid:
@@ -80,7 +89,7 @@ def cli(account=None, role=None, profile='default'):
 
         check_for_aws_dir()
         write_to_cred_file(aws_file('creds'), aws_session_token, profile)
-        write_to_config_file(aws_file('config'), net_id, 'us-west-2', profile)
+        write_to_config_file(aws_file('config'), net_id, 'us-west-2', profile, account_role.role_name, account_role.account_name)
     
         if account_role.role_name == "AccountAdministrator":
             print("Now logged into {}{}{}@{}{}{}".format(Colors.red,account_role.role_name, Colors.white, Colors.yellow,account_role.account_name,Colors.normal))
@@ -124,10 +133,15 @@ def write_to_cred_file(file, aws_session_token, profile):
         config.write(configfile)
 
 
-def write_to_config_file(file, net_id, region, profile):
+def write_to_config_file(file, net_id, region, profile, role, account):
+    one_hour = datetime.timedelta(hours=1)
+    expires = datetime.datetime.now() + one_hour
     config = open_config_file(file)
     config[profile] = {
-        'region': region, 'adfs_netid': net_id
+        'region': region,
+        'adfs_netid': net_id,
+        'adfs_role': f'{role}@{account}',
+        'adfs_expires': expires.strftime('%m-%d-%Y %H:%M')
     }
     with open(file, 'w') as configfile:
         config.write(configfile)
@@ -141,5 +155,44 @@ def load_last_netid(file, profile):
         return ''
 
 
+def get_status_message(config, profile):
+    if config.has_option(profile, 'adfs_role') and config.has_option(profile, 'adfs_expires'):
+        expires = check_expired(config[profile]['adfs_expires'])
+        account_name = f"{Colors.cyan}{config[profile]['adfs_role']}"
+        if expires == 'Expired':
+            expires_msg = f"{Colors.red}{expires} at: {config[profile]['adfs_expires']}"
+        else:
+            expires_msg = f"{Colors.yellow}{expires} at: {config[profile]['adfs_expires']}"
+        return f"{account_name} {Colors.white}- {expires_msg}"
+    else:
+        return f"{Colors.red}Couldn't find status info"
+
+
+def get_status(file, profile='default'):
+    config = open_config_file(file)
+    if profile == 'all':
+        for x in config:
+            if x == 'DEFAULT':
+                continue
+            message = get_status_message(config, x)
+            print(f"{Colors.white}{x} - {message}")
+        return
+    else:
+        if config.has_section(profile):
+            message = get_status_message(config, profile)
+            print(message)
+        else:
+            print(f"{Colors.red}Couldn't find profile: {profile}")
+        return
+
+
+def check_expired(expires):
+    expires = datetime.datetime.strptime(expires, '%m-%d-%Y %H:%M')
+    if expires > datetime.datetime.now():
+        return 'Expires'
+    else:
+        return 'Expired'
+
+
 if __name__ == '__main__':
     main()
diff --git a/setup.py b/setup.py
index a19a657..bc54bed 100644
--- a/setup.py
+++ b/setup.py
@@ -4,6 +4,7 @@ import os
 import sys
 import subprocess
 from setuptools import setup, find_packages
+from byu_awslogin import index
 
 if not sys.version.startswith('3.6'):
     sys.stderr.write("byu_awslogin requires python 3.6\n")
@@ -17,9 +18,7 @@ def get_requirements():
         lines = [dep for dep in obj.read().split('\n') if dep]
         return lines
 
-VERSION = ''
-with open('VERSION') as verobj:
-    VERSION = verobj.read()
+VERSION = index.__VERSION__
 
 setup(name='byu_awslogin',
       version=VERSION,


Test patch:
diff --git a/test/byu_awslogin/test_index.py b/test/byu_awslogin/test_index.py
index 2a8be70..bc940fd 100644
--- a/test/byu_awslogin/test_index.py
+++ b/test/byu_awslogin/test_index.py
@@ -1,6 +1,7 @@
 import pytest
 import os.path
 import configparser
+import datetime
 from byu_awslogin import index
 
 
@@ -25,14 +26,21 @@ def read_config_file(file):
 def profile():
     return 'default'
 
+
 @pytest.fixture
 def fake_config_file(tmpdir):
     dir = tmpdir.dirpath(".aws/")
     file = dir + 'config'
+    exp_date = datetime.datetime.now() + datetime.timedelta(hours=1)
     if not os.path.exists(dir):
         os.makedirs(dir)
     config = configparser.ConfigParser()
-    config['default'] = {'region': 'fake-region-1', 'adfs_netid': 'fake_netid'}
+    config['default'] = {
+        'region': 'fake-region-1',
+        'adfs_netid': 'fake_netid',
+        'adfs_role': 'FakeRole@FakeAccount',
+        'adfs_expires': exp_date.strftime('%m-%d-%Y %H:%M')
+    }
     with open(file, 'w') as write_file:
         config.write(write_file)
     return file
@@ -65,15 +73,43 @@ def test_write_to_cred_file(aws_cred_file, profile):
 def test_write_to_config_file(aws_config_file, profile):
     net_id = 'fake_netid'
     region = 'fakeRegion'
-    index.write_to_config_file(aws_config_file, net_id, region, profile)
+    role = 'FakeRole'
+    account = 'FakeAccount'
+    index.write_to_config_file(aws_config_file, net_id, region, profile, role, account)
     config = read_config_file(aws_config_file)
-    assert config['default'] == {'region': 'fakeRegion', 'adfs_netid': 'fake_netid'}
+    assert config.has_section('default')
+    assert config.has_option('default', 'region')
+    assert config.has_option('default', 'adfs_netid')
+    assert config.has_option('default', 'adfs_role')
+    assert config.has_option('default', 'adfs_expires')
 
 
 def test_load_last_netid(fake_config_file, profile):
     assert index.load_last_netid(str(fake_config_file), profile) == 'fake_netid'
 
 
-@pytest.mark.skip(reason="not sure how to test this yet")
-def test_cli():
-    pass
+def test_cli_version(capfd):
+    index.cli(version=True)
+    out, err = capfd.readouterr()
+    assert out.strip() == f'awslogin version: {index.__VERSION__}'
+
+
+def test_get_status(fake_config_file, profile, capfd):
+    index.get_status(str(fake_config_file), profile)
+    out, err = capfd.readouterr()
+    assert 'FakeRole@FakeAccount' in out
+
+
+def test_get_status_fail(fake_config_file, capfd):
+    index.get_status(str(fake_config_file), 'fake')
+    out, err = capfd.readouterr()
+    assert "Couldn't find profile:" in out
+
+
+def test_check_expired():
+    expired = datetime.datetime.now() - datetime.timedelta(hours=1)
+    result = index.check_expired(expired.strftime('%m-%d-%Y %H:%M'))
+    assert result == 'Expired'
+    expires = datetime.datetime.now() + datetime.timedelta(hours=1)
+    result = index.check_expired(expires.strftime('%m-%d-%Y %H:%M'))
+    assert result == 'Expires'