From 5344213c7cf4ff0d852f3c3afc0d6c48fdd4acd0 Mon Sep 17 00:00:00 2001 From: Snigdhajyoti Ghosh Date: Fri, 24 Nov 2023 02:43:14 +0530 Subject: [PATCH] Add init command to create alias. --- aws_fusion/__init__.py | 2 +- aws_fusion/app.py | 10 +++++-- aws_fusion/aws/session.py | 4 +++ aws_fusion/commands/init.py | 41 ++++++++++++++++++++++++++++ aws_fusion/okta/api.py | 3 +-- setup.py | 54 ------------------------------------- 6 files changed, 55 insertions(+), 59 deletions(-) create mode 100644 aws_fusion/commands/init.py diff --git a/aws_fusion/__init__.py b/aws_fusion/__init__.py index 64477cf..6f4fa58 100644 --- a/aws_fusion/__init__.py +++ b/aws_fusion/__init__.py @@ -1 +1 @@ -__version__ = '1.2' +__version__ = '1.3' diff --git a/aws_fusion/app.py b/aws_fusion/app.py index a3ed5b9..8332aca 100644 --- a/aws_fusion/app.py +++ b/aws_fusion/app.py @@ -1,12 +1,13 @@ import argparse +import logging from importlib.metadata import version -from .commands import open_browser, iam_user_credentials, generate_okta_device_auth_credentials - +from .commands import open_browser, iam_user_credentials, generate_okta_device_auth_credentials, init def main(): global_parser = argparse.ArgumentParser(add_help=False) global_parser.add_argument('-v', '--version', action='version', help="Display the version of this tool", version=version("aws_fusion")) + global_parser.add_argument('--debug', action='store_true', help='Turn on debug logging') main_parser = argparse.ArgumentParser(prog='aws-fusion', description='Unified CLI tool for streamlined AWS operations, enhancing developer productivity', parents=[global_parser]) subparsers = main_parser.add_subparsers(dest='command', required=True, help='Available commands') @@ -14,8 +15,13 @@ def main(): open_browser.setup(subparsers, global_parser) iam_user_credentials.setup(subparsers, global_parser) generate_okta_device_auth_credentials.setup(subparsers, global_parser) + init.setup(subparsers, global_parser) args = main_parser.parse_args() + + if args.debug: + logging.basicConfig(level=logging.DEBUG) + # Call the associated function for the selected sub-command if hasattr(args, 'func'): args.func(args) diff --git a/aws_fusion/aws/session.py b/aws_fusion/aws/session.py index 032d97b..92b7d29 100644 --- a/aws_fusion/aws/session.py +++ b/aws_fusion/aws/session.py @@ -1,7 +1,9 @@ import boto3 +import logging from botocore.utils import JSONFileCache +LOG = logging.getLogger(__name__) class TokenGenerationException(Exception): """Exception for credential not having token""" @@ -9,12 +11,14 @@ class TokenGenerationException(Exception): def credentials(profile_name, region_name): + LOG.debug(f'Creating boto3 session for profile {profile_name} and region {region_name}') session = boto3.Session(profile_name=profile_name, region_name=region_name) __update_credential_provider_cache(session) creds = session.get_credentials() if creds.token is None: + LOG.error("No session credential found") raise TokenGenerationException() return creds, session.region_name diff --git a/aws_fusion/commands/init.py b/aws_fusion/commands/init.py new file mode 100644 index 0000000..2026f22 --- /dev/null +++ b/aws_fusion/commands/init.py @@ -0,0 +1,41 @@ +import os +import configparser +import sys +import logging + +LOG = logging.getLogger(__name__) + + +def setup(subparsers, parent_parser): + summary = 'Initilize fusion app with creation of aws fusion alias' + parser = subparsers.add_parser('init', description=summary, help=summary, parents=[parent_parser]) + parser.set_defaults(func=run) + +def run(args): + cli_dir = os.path.expanduser(os.path.join('~', '.aws', 'cli')) + + def create_alias(config: configparser.ConfigParser): + LOG.debug('Creating fusion aws alias') + if not config.has_section('toplevel'): + config.add_section('toplevel') + config['toplevel']['fusion'] = f'!{sys.executable} -m aws_fusion.app' + + def update_aws_cli_alias_file(): + if not os.path.isdir(cli_dir): + LOG.debug(f"Path {cli_dir} doesn't exists, creating" ) + os.makedirs(cli_dir) + cli_alias_full_path = os.path.join(cli_dir, 'alias') + config = configparser.ConfigParser() + + if os.path.isfile(cli_alias_full_path): + LOG.debug(f'Found alias file {cli_alias_full_path}') + config.read(cli_alias_full_path) + + create_alias(config) + + with os.fdopen(os.open(cli_alias_full_path, os.O_WRONLY | os.O_CREAT, 0o600), 'w') as f: + f.truncate() + config.write(f) + LOG.debug(f'Updated alias file {cli_alias_full_path}') + + update_aws_cli_alias_file() diff --git a/aws_fusion/okta/api.py b/aws_fusion/okta/api.py index a5576dc..4258460 100644 --- a/aws_fusion/okta/api.py +++ b/aws_fusion/okta/api.py @@ -97,6 +97,5 @@ def saml_assertion(org_domain, session_token): {"Name": "https://aws.amazon.com/SAML/Attributes/SessionDuration"}) session_duration = session_duration_xml.find("saml2:AttributeValue").text - LOG.debug('Got valid SAML response') - time.sleep(5) + LOG.debug(f'Got valid SAML response: {saml_response}') return saml_response, roles, int(session_duration) diff --git a/setup.py b/setup.py index f98028a..4e6981a 100755 --- a/setup.py +++ b/setup.py @@ -2,10 +2,6 @@ import re import os from setuptools import setup, find_packages -from setuptools.command.develop import develop -from setuptools.command.install_scripts import install_scripts -import configparser -import sys here = os.path.abspath(os.path.dirname(__file__)) @@ -23,52 +19,6 @@ def find_version(*file_paths): raise RuntimeError("Unable to find version string.") -def update_aws_cli_alias(command_subclass): - """ - A decorator for classes subclassing one of the setuptools commands. - It modifies the run() method so that it prints a friendly greeting. - """ - orig_run = command_subclass.run - cli_dir = os.path.expanduser(os.path.join('~', '.aws', 'cli')) - - def create_alias(config: configparser.ConfigParser): - if not config.has_section('toplevel'): - config.add_section('toplevel') - config['toplevel']['fusion'] = f'!{sys.executable} -m aws_fusion.app' - - def update_aws_cli_alias_file(): - if not os.path.isdir(cli_dir): - os.makedirs(cli_dir) - cli_alias_full_path = os.path.join(cli_dir, 'alias') - config = configparser.ConfigParser() - - if os.path.isfile(cli_alias_full_path): - config.read(cli_alias_full_path) - - create_alias(config) - - with os.fdopen(os.open(cli_alias_full_path, os.O_WRONLY | os.O_CREAT, 0o600), 'w') as f: - f.truncate() - config.write(f) - - def modified_run(self): - orig_run(self) - update_aws_cli_alias_file() - - command_subclass.run = modified_run - return command_subclass - - -@update_aws_cli_alias -class CustomDevelopCommand(develop): - pass - - -@update_aws_cli_alias -class CustomInstallScriptsCommand(install_scripts): - pass - - setup( name='aws-fusion', version=find_version('aws_fusion', '__init__.py'), @@ -84,10 +34,6 @@ class CustomInstallScriptsCommand(install_scripts): long_description=read("README.md"), long_description_content_type='text/markdown', packages=find_packages(), - cmdclass={ - 'develop': CustomDevelopCommand, - 'install_scripts': CustomInstallScriptsCommand - }, entry_points={ 'console_scripts': [ 'aws-fusion = aws_fusion.app:main',