Merge pull request #19 from aitomatic/feat/change_credential_format

Change format of credential and .aito files from json to ini

src/aito.py

@@ -1,16 +1,17 @@
 import click
-import json
 from src.execute.main import execute
 from src.deploy.main import deploy
 from src.login.main import login
 from src.logout.main import logout
 from src.run.main import run
-from src.constants import CREDENTIAL_FILE
+from src.constants import CREDENTIAL_FILE, AITOMATIC_PROFILE
+from src.utils import read_ini_file
 
 
 def load_config():
     if CREDENTIAL_FILE.exists():
-        return json.loads(CREDENTIAL_FILE.read_text())
+        credentials = read_ini_file(CREDENTIAL_FILE)
+        return credentials[AITOMATIC_PROFILE]
     else:
         return {}
 

src/constants.py

@@ -1,4 +1,6 @@
+from os import getenv
 from pathlib import Path
 
 
 CREDENTIAL_FILE = Path.home().joinpath('.aitomatic/credentials')
+AITOMATIC_PROFILE = getenv('AITOMATIC_PROFILE', 'default')

src/login/main.py

@@ -1,14 +1,14 @@
 import click
+from configparser import ConfigParser
 import requests
 from multiprocessing import Process, Queue
 from http.server import HTTPServer
 import time
-import json
 from functools import update_wrapper, partial
 
 from .server import LoginServer
 from src.utils import get_random_string, create_code_challenger, create_code_verifier
-from src.constants import CREDENTIAL_FILE
+from src.constants import CREDENTIAL_FILE, AITOMATIC_PROFILE
 from src.logout.main import remove_local_credential_file
 
 ORG = 'aitomaticinc.us.auth0.com'
@@ -141,11 +141,9 @@ def wait_for_login_callback(obj):
 
             if polling_data.get('access_token') is not None:
                 save_credential(
-                    {
-                        'access_token': polling_data['access_token'],
-                        'refresh_token': polling_data.get('refresh_token', ''),
-                        'id': polling_data.get('id_token', ''),
-                    }
+                    access_token=polling_data['access_token'],
+                    refresh_token=polling_data.get('refresh_token', ''),
+                    id_token=polling_data.get('id_token', ''),
                 )
                 click.echo('Login successfully')
                 exit(0)
@@ -159,13 +157,22 @@ def wait_for_login_callback(obj):
 
 
 @click.pass_obj
-def save_credential(obj, data):
-    obj['access_token'] = data['access_token']
-    obj['refresh_token'] = data['refresh_token']
-    obj['id'] = data['id']
+def save_credential(obj, access_token, refresh_token, id_token):
+    obj['access_token'] = access_token
+    obj['refresh_token'] = refresh_token
+    obj['id_token'] = id_token
+
+    parser = ConfigParser()
+    parser[AITOMATIC_PROFILE] = {
+        'access_token': access_token,
+        'refresh_token': refresh_token,
+        'id_token': id_token,
+    }
+
     if not CREDENTIAL_FILE.exists():
         CREDENTIAL_FILE.parent.mkdir(parents=True, exist_ok=True)
-    CREDENTIAL_FILE.write_text(json.dumps(data))
+    with open(CREDENTIAL_FILE, 'w') as f:
+        parser.write(f)
 
 
 def authenticated(f):
@@ -217,11 +224,9 @@ def refresh_token(obj):
     if res.status_code == 200:
         data = res.json()
         save_credential(
-            {
-                'access_token': data['access_token'],
-                'refresh_token': token,
-                'id': data.get('id_token', ''),
-            }
+            access_token=data['access_token'],
+            refresh_token=token,
+            id_token=data.get('id_token', ''),
         )
     else:
         prompt_login()

src/logout/main.py

@@ -4,8 +4,7 @@
 
 
 @click.command()
-@click.pass_obj
-def logout(obj):
+def logout():
     '''Logout from Aitomatic cloud'''
     remove_local_credential_file()
     click.echo('Logout successfully. Please login to use other commands.')

src/run/main.py

@@ -1,25 +1,26 @@
 import click
 import json
-from configparser import ConfigParser
 from pathlib import Path
 from src.login.main import authenticated
 from src.execute.app import execute_app
+from src.utils import read_ini_file
+from src.constants import AITOMATIC_PROFILE
 
 
 @click.command()
 @click.option(
     '-c',
     '--config',
-    'config_file',
+    'app_config_file',
     type=click.STRING,
     help='ini file to run the app',
 )
 @authenticated
-def run(config_file):
+def run(app_config_file):
     '''Run the app based on .aito config file'''
     aito_config = AitoConfig()
-    if config_file is not None:
-        aito_config.set_app_config(config_file)
+    if app_config_file is not None:
+        aito_config.set_app_config(app_config_file)
 
     data = execute_app(app_name=aito_config.app_name, data=aito_config.app_config)
     click.echo(data)
@@ -44,27 +45,17 @@ def read_aito_file(self):
             exit(1)
 
         try:
-            return json.loads(config_files[0].read_text())
-        except json.decoder.JSONDecodeError:
-            click.echo("Can't read .aito config file.")
+            return read_ini_file(config_files[0])[AITOMATIC_PROFILE]
+        except KeyError:
+            click.echo(
+                f"Can't read .aito config file with profile {AITOMATIC_PROFILE}."
+            )
             exit(1)
 
-    def convert_ini_config_to_dict(self, config_content):
-        parser = ConfigParser()
-        result = {}
-
-        parser.read_string(config_content)
-        for section in parser.sections():
-            result[section] = {}
-            for name, value in parser.items(section):
-                result[section][name] = value
-
-        return result
-
-    def set_app_config(self, config_file):
+    def set_app_config(self, app_config_file):
         try:
-            file_path = Path.cwd().joinpath(config_file)
-            self.app_config = self.convert_ini_config_to_dict(file_path.read_text())
+            file_path = Path.cwd().joinpath(app_config_file)
+            self.app_config = read_ini_file(file_path)
         except FileNotFoundError:
             click.echo("Can't read app config file")
             exit(1)

src/utils.py

@@ -1,13 +1,15 @@
-from os import urandom
 from base64 import urlsafe_b64encode
+from configparser import ConfigParser
+from os import urandom
 from hashlib import sha256
+from pathlib import Path
 from re import sub
 
 import string
 import random
 
 
-def get_random_string(length):
+def get_random_string(length: int) -> str:
     letters = string.ascii_lowercase
     return ''.join(random.choice(letters) for i in range(length))
 
@@ -19,8 +21,21 @@ def create_code_verifier():
     return code_verifier
 
 
-def create_code_challenger(code_verifier):
+def create_code_challenger(code_verifier: str) -> str:
     code_challenge = sha256(code_verifier.encode('utf-8')).digest()
     code_challenge = urlsafe_b64encode(code_challenge).decode('utf-8')
     code_challenge = code_challenge.replace('=', '')
     return code_challenge
+
+
+def read_ini_file(file_path: Path) -> dict:
+    parser = ConfigParser()
+    result = {}
+
+    parser.read_string(file_path.read_text())
+    for section in parser.sections():
+        result[section] = {}
+        for name, value in parser.items(section):
+            result[section][name] = value
+
+    return result

version.txt

@@ -1 +1 @@
-0.8.0
+0.9.0