estela-cli: Add command to save different contexts

.gitignore

@@ -12,3 +12,4 @@ venv/
 .venv/
 .DS_Store
 site/
+.vscode/

estela_cli/context.py

@@ -1,15 +1,21 @@
+import yaml
 import click
 import requests
 
+from getpass import getpass
 from estela_cli.login import env_login, yaml_login
+from estela_cli.estela_client import EstelaClient
 from estela_cli.utils import (
     get_estela_auth,
+    update_estela_auth,
     get_estela_settings,
+    get_estela_config,
     get_host_from_env,
     get_username_from_env,
     get_password_from_env,
+    get_estela_config_path,
 )
-from estela_cli.templates import ESTELA_AUTH_NAME, OK_EMOJI, BAD_EMOJI
+from estela_cli.templates import ESTELA_AUTH_NAME, ESTELA_CONFIG_NAME, OK_EMOJI, BAD_EMOJI
 
 
 SHORT_HELP = "Show your current context"
@@ -30,12 +36,73 @@ def test_host(host):
     assert "projects" in response.json()
 
 
+def prompt_context(name, username=None, password=None, host=None):
+    try:
+        if host is None:
+            host = click.prompt("Host")
+        if username is None:
+            username = click.prompt("Username")
+        if password is None:
+            password = getpass()
+
+        estela_config_path = get_estela_config_path()
+        estela_config = get_estela_config()
+
+        if estela_config is None:
+            estela_config = {}
+
+        estela_config[name] = {
+            "host": host,
+            "username": username,
+            "password": password,
+        }
+        with open(estela_config_path, "w") as estela_config_yaml:
+            yaml.dump(estela_config, estela_config_yaml)
+            click.echo(
+                "Successful login. Context {} stored in ~/{}.".format(
+                    name, ESTELA_CONFIG_NAME
+                )
+            )
+        estela_client = EstelaClient(host, username, password)
+
+    except:
+        raise Exception("Unable to login with provided credentials.")
+
+    return estela_client
+
+
 @click.command(name="context", short_help=SHORT_HELP)
-def estela_command():
+@click.argument("name", required=False)
+def estela_command(name):
     host = get_host_from_env()
     username = get_username_from_env()
     password = get_password_from_env()
 
+    if name:
+        click.echo(f"Checking your current context for {name}...")
+        estela_config = get_estela_config()
+
+        if estela_config and estela_config.get(name, None):
+            click.echo("Context {} found.".format(name))
+            try:
+                test_host(estela_config[name]["host"])
+                click.echo(OK_HOST.format(estela_config[name]["host"]))
+                estela_client = EstelaClient(
+                    estela_config[name]["host"],
+                    estela_config[name]["username"],
+                    estela_config[name]["password"],
+                )
+            except:
+                click.echo(BAD_HOST)
+                return
+        else:
+            click.echo(
+                "Context {} not found.\nInitializing context...".format(name)
+            )
+            estela_client = prompt_context(name, username, password, host)
+        update_estela_auth(estela_client.host, estela_client.token)
+        return
+
     if host is None or username is None or password is None:
         estela_auth = get_estela_auth()
 

estela_cli/login.py

@@ -2,16 +2,16 @@
 import click
 
 from getpass import getpass
-from string import Template
 from estela_cli.estela_client import EstelaClient
 from estela_cli.utils import (
     get_host_from_env,
     get_username_from_env,
     get_password_from_env,
     get_estela_auth,
-    get_home_path,
+    update_estela_auth,
+    get_estela_auth_path,
 )
-from estela_cli.templates import ESTELA_AUTH_NAME, ESTELA_AUTH
+from estela_cli.templates import ESTELA_AUTH_NAME
 
 
 DEFAULT_ESTELA_API_HOST = "http://localhost"
@@ -109,8 +109,7 @@ def login(username=None, password=None, host=None):
     help="API endpoint to send the requests. If host is not given, it will be asked",
 )
 def estela_command(username, password, host):
-    home_path = get_home_path()
-    estela_auth_path = os.path.join(home_path, ESTELA_AUTH_NAME)
+    estela_auth_path = get_estela_auth_path()
 
     if os.path.exists(estela_auth_path):
         raise click.ClickException(
@@ -122,15 +121,4 @@ def estela_command(username, password, host):
     except Exception as ex:
         raise click.ClickException(str(ex))
 
-    template = Template(ESTELA_AUTH)
-    values = {
-        "estela_host": estela_client.host,
-        "estela_token": estela_client.token,
-    }
-    result = template.substitute(values)
-
-    with open(estela_auth_path, "w") as estela_auth_yaml:
-        estela_auth_yaml.write(result)
-        click.echo(
-            "Successful login. API Token stored in ~/{}.".format(ESTELA_AUTH_NAME)
-        )
+    update_estela_auth(estela_client.host, estela_client.token)

estela_cli/templates.py

@@ -1,6 +1,7 @@
 # Auth templates
 
 ESTELA_AUTH_NAME = ".estela.yaml"
+ESTELA_CONFIG_NAME = ".estela-config.yaml"
 
 ESTELA_AUTH = """\
 token: $estela_token

estela_cli/utils.py

@@ -4,10 +4,13 @@
 import json
 import click
 
+from string import Template
 from datetime import datetime
 from estela_cli.templates import (
+    ESTELA_CONFIG_NAME,
     ESTELA_AUTH_NAME,
     ESTELA_YAML_NAME,
+    ESTELA_AUTH,
     ESTELA_DIR,
     DATA_DIR,
     DOCKERFILE_NAME,
@@ -55,9 +58,14 @@ def get_estela_settings():
     return estela_config
 
 
-def get_estela_auth():
+def get_estela_auth_path():
     home_path = get_home_path()
     estela_auth_path = os.path.join(home_path, ESTELA_AUTH_NAME)
+    return estela_auth_path
+
+
+def get_estela_auth():
+    estela_auth_path = get_estela_auth_path()
 
     if not os.path.exists(estela_auth_path):
         return None
@@ -68,6 +76,42 @@ def get_estela_auth():
     return estela_auth
 
 
+def update_estela_auth(host, token):
+    estela_auth_path = get_estela_auth_path()
+    template = Template(ESTELA_AUTH)
+    values = {
+        "estela_host": host,
+        "estela_token": token,
+    }
+    result = template.substitute(values)
+
+    with open(estela_auth_path, "w") as estela_auth_yaml:
+        estela_auth_yaml.write(result)
+        click.echo(
+            "Successful login. API Token stored in ~/{}.".format(ESTELA_AUTH_NAME)
+        )
+
+
+def get_estela_config_path():
+    home_path = get_home_path()
+    estela_config_path = os.path.join(home_path, ESTELA_CONFIG_NAME)
+    return estela_config_path
+
+
+def get_estela_config():
+    estela_config_path = get_estela_config_path()
+
+    if not os.path.exists(estela_config_path):
+        file = open(estela_config_path, "x")
+        file.close()
+        return {}
+
+    with open(estela_config_path, "r") as estela_config_yaml:
+        estela_config = yaml.full_load(estela_config_yaml)
+
+    return estela_config
+
+
 def format_time(date):
     date = datetime.strptime(date, "%Y-%m-%dT%H:%M:%S.%fZ")
     return date.strftime("%Y-%m-%d %H:%M")