add aws cli alias for both console & credential-process-from-system

setup.py

@@ -2,6 +2,9 @@
 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
 
 here = os.path.abspath(os.path.dirname(__file__))
 
@@ -18,12 +21,56 @@ def find_version(*file_paths):
         return version_match.group(1)
     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']['console'] = '!aws-console'
+        config['toplevel']['credential-process-from-system'] = '!aws-credential-process-from-system'
+
+    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_console',
     version=find_version('aws_console', '__init__.py'),
     description='AWS Console Login Utility',
-    keywords = [
+    keywords=[
         'aws',
         'aws-sdk',
         'aws-cli',
@@ -33,6 +80,10 @@ def find_version(*file_paths):
     ],
     long_description=read("README.md"),
     packages=find_packages(),
+    cmdclass={
+        'develop': CustomDevelopCommand,
+        'install_scripts': CustomInstallScriptsCommand
+    },
     entry_points={
         'console_scripts': [
             'aws-console = aws_console:aws_console',