Merge pull request #15 from aitomatic/feat/refresh_key_flow

Add flow to refresh token automatically when it expired

Author: khoama

DEVELOPING.md

@@ -5,83 +5,83 @@
 1. python3
 2. pip
 3. virtualenv
-```shell
-pip install virtualenv
-```
+    ```shell
+    pip install virtualenv
+    ```
 
 ## Steps to setup dev environment
 
 1. Create a python3 virtual environment
-```shell
-virtualenv .venv -p python3
-```
+    ```shell
+    virtualenv .venv -p python3
+    ```
 2. Activate the virtual environment created in step 1
-```shell
-source .venv/bin/activate
-```
+    ```shell
+    source .venv/bin/activate
+    ```
 3. Install `setuptools` to package code and `black` to format python code
-```shell
-pip install setuptools black
-```
+    ```shell
+    pip install setuptools black
+    ```
 4. Install libraries in file requirements.txt
-```shell
-pip install -r requirements.txt
-```
+    ```shell
+    pip install -r requirements.txt
+    ```
 
 ## Steps to test code in local
 
 1. Run `black` command to format code
-```shell
-black -S .
-```
+    ```shell
+    black -S .
+    ```
 2. Run command to test script
-```shell
-python src/aito.py
-```
+    ```shell
+    python src/aito.py
+    ```
 
 ## Steps to install the CLI to virtual environment and test it
 
 1. Run command to install CLI
-```shell
-pip install -e .
-```
+    ```shell
+    pip install -e .
+    ```
 2. Verify the CLI
-```shell
-which aito
-```
+    ```shell
+    which aito
+    ```
 3. Run `aito` command
-```shell
-aito
-```
+    ```shell
+    aito
+    ```
 
 ## Steps to package and distribute CLI to TestPyPI
 
 1. Install `build` to generate distribution packages and `twine` to distribute package to PyPI
-```shell
-pip install build twine
-```
+    ```shell
+    pip install build twine
+    ```
 2. Run `build` command as the root folder, where file `pyproject.toml` is located
-```shell
-rm -rf dist
-python -m build
-```
-After that command, we will have `tar.gz` and `.whl` files in `dist` folder
+    ```shell
+    rm -rf dist
+    python -m build
+    ```
+    After that command, we will have `tar.gz` and `.whl` files in `dist` folder
 
 3. Register an account in TestPyPI and create an API token with `Entire account` scope
 4. Using twine to upload the distribution packages created in step 2 to TestPyPI
-```shell
-twine upload --repository testpypi --skip-existing dist/*
-```
-`--repository` used to choose upload to PyPI or TestPyPI, `--skip-existing` if we want to distribute further versions of the cli.
+    ```shell
+    twine upload --repository testpypi --skip-existing dist/*
+    ```
+    `--repository` used to choose upload to PyPI or TestPyPI, `--skip-existing` if we want to distribute further versions of the cli.
 
-You will be prompted for a username and password. For the username, use __token__. For the password, use the token value, including the pypi- prefix.
+    You will be prompted for a username and password. For the username, use __token__. For the password, use the token value, including the pypi- prefix.
 5. Using another virtual environment and install the `aitomatic-cli` using pip to verify that it works
-```shell
-deactivate
-virtualenv .venv-test -p python3
-source .venv-test/bin/activate
-pip install -i https://test.pypi.org/simple/ aitomatic
-```
+    ```shell
+    deactivate
+    virtualenv .venv-test -p python3
+    source .venv-test/bin/activate
+    pip install -i https://test.pypi.org/simple/ aitomatic
+    ```
 
 ## Steps to package and distribute CLI to PyPI
 
@@ -92,13 +92,13 @@ Similar to steps to distribute to TestPyPI, except:
 ## Temporary: If the deploy command failed due to change in service api. Update api base:
 
 1. Clone https://github.com/aitomatic/ai-cloud
-2. Follow instruction and select dev stack
+2. Follow instruction in `infrastructure/DEVELOPING.md` and select `dev` stack
 3. Get the service host name
-```shell
-pulumi stack select dev
-pulumi stack output kodaServiceHostname
-```
+    ```shell
+    pulumi stack select dev
+    pulumi stack output kodaServiceHostname
+    ```
 4. Set AI_CLI_API_BASE env var to 
-```shell
-export AI_CLI_API_BASE = "http://" + <hostname from step 3>
-```
+    ```shell
+    export AI_CLI_API_BASE = "http://" + <hostname from step 3>
+    ```

src/login/main.py

@@ -107,6 +107,7 @@ def initiate_login_flow(obj):
         f'&scope={SCOPE}'
         f'&audience={AUDIENCE}'
         f'&state={login_seed}'
+        f'&max_age=3600'
     )
 
     click.launch(url)
@@ -145,7 +146,7 @@ def wait_for_login_callback(obj):
                         'id': polling_data.get('id_token', ''),
                     }
                 )
-                click.echo('Login successful')
+                click.echo('Login successfully')
                 exit(0)
             return
 
@@ -169,11 +170,10 @@ def save_credential(obj, data):
 def authenticated(f):
     @click.pass_obj
     def wrapper(obj, *args, **kwargs):
-        token = obj and obj.get('access_token')
+        token = obj.get('access_token')
 
-        if token is None:
+        if token is None or len(token) == 0:
             prompt_login()
-            exit(1)
 
         res = requests.get(
             url='https://{}/userinfo'.format(ORG),
@@ -182,13 +182,44 @@ def wrapper(obj, *args, **kwargs):
 
         if res.status_code == 200:
             f(*args, **kwargs)
+        elif res.status_code == 401:
+            refresh_token()
+            f(*args, **kwargs)
         else:
-            click.echo(res.text)
             prompt_login()
-            exit(1)
 
     return update_wrapper(wrapper, f)
 
 
 def prompt_login():
     click.echo("You're not logged in. Please run `aito login` first.")
+    exit(1)
+
+
+@click.pass_obj
+def refresh_token(obj):
+    token = obj.get('refresh_token')
+
+    if token is None or len(token) == 0:
+        prompt_login()
+
+    res = requests.post(
+        url='https://{}/oauth/token'.format(ORG),
+        data={
+            'client_id': CLIENT_ID,
+            'grant_type': 'refresh_token',
+            'refresh_token': obj['refresh_token'],
+        },
+        headers={'content-type': 'application/x-www-form-urlencoded'},
+    )
+    if res.status_code == 200:
+        data = res.json()
+        save_credential(
+            {
+                'access_token': data['access_token'],
+                'refresh_token': token,
+                'id': data.get('id_token', ''),
+            }
+        )
+    else:
+        prompt_login()

version.txt

@@ -1 +1 @@
-0.6.0
+0.6.1