From 899ec8313c8ccabeb2e10eb841c41a6c6a9eaaf4 Mon Sep 17 00:00:00 2001
From: Snigdhajyoti Ghosh <snigdhasjg@users.noreply.github.com>
Date: Mon, 13 Nov 2023 01:44:33 +0530
Subject: [PATCH] add aws cli alias for both console &
 credential-process-from-system

---
 setup.py | 36 +++++++++++++++++++++++++++++++++++-
 1 file changed, 35 insertions(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index 4a3ea70..3e229ab 100755
--- a/setup.py
+++ b/setup.py
@@ -2,6 +2,8 @@
 import re
 import os
 from setuptools import setup, find_packages
+from setuptools.command.install import install
+import configparser
 
 here = os.path.abspath(os.path.dirname(__file__))
 
@@ -19,11 +21,40 @@ def find_version(*file_paths):
     raise RuntimeError("Unable to find version string.")
 
 
+class CustomInstall(install):
+    CLI_DIR = os.path.expanduser(os.path.join('~', '.aws', 'cli'))
+
+    def run(self):
+        self.update_aws_cli_alias_file()
+        return super().run()
+
+    def update_aws_cli_alias_file(self):
+        if not os.path.isdir(self.CLI_DIR):
+            os.makedirs(self.CLI_DIR)
+        cli_alias_full_path = os.path.join(self.CLI_DIR, 'alias')
+        config = configparser.ConfigParser()
+
+        if os.path.isfile(cli_alias_full_path):
+            config.read(cli_alias_full_path)
+
+        self.create_console_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 create_console_alias(self, 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'
+
+
 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 +64,9 @@ def find_version(*file_paths):
     ],
     long_description=read("README.md"),
     packages=find_packages(),
+    cmdclass={
+        'install': CustomInstall
+    },
     entry_points={
         'console_scripts': [
             'aws-console = aws_console:aws_console',