From 8c0b1ea6301e5cce66f750b3feb52c5341dd9946 Mon Sep 17 00:00:00 2001
From: Taha Yassine Kraiem <tahayk2@gmail.com>
Date: Tue, 30 Nov 2021 16:48:23 +0100
Subject: [PATCH] feat(api): EE changed credentials endpoints

---
 ee/api/app.py                                 |  6 ++---
 ee/api/chalicelib/blueprints/app/v1_api_ee.py | 14 +++++++++++
 ee/api/chalicelib/blueprints/bp_ee.py         | 24 ++++---------------
 ee/api/chalicelib/utils/assist_helper.py      | 19 +++++++++++++++
 4 files changed, 40 insertions(+), 23 deletions(-)
 create mode 100644 ee/api/chalicelib/blueprints/app/v1_api_ee.py
 create mode 100644 ee/api/chalicelib/utils/assist_helper.py

diff --git a/ee/api/app.py b/ee/api/app.py
index cc901c6f8..e12b64e0b 100644
--- a/ee/api/app.py
+++ b/ee/api/app.py
@@ -5,15 +5,14 @@ from sentry_sdk import configure_scope
 from chalicelib import _overrides
 from chalicelib.blueprints import bp_authorizers
 from chalicelib.blueprints import bp_core, bp_core_crons
-from chalicelib.blueprints.app import v1_api
 from chalicelib.blueprints import bp_core_dynamic, bp_core_dynamic_crons
+from chalicelib.blueprints import bp_ee, bp_ee_crons, bp_saml
+from chalicelib.blueprints.app import v1_api, v1_api_ee
 from chalicelib.blueprints.subs import bp_dashboard
 from chalicelib.utils import helper
 from chalicelib.utils import pg_client
 from chalicelib.utils.helper import environ
 
-from chalicelib.blueprints import bp_ee, bp_ee_crons, bp_saml
-
 app = Chalice(app_name='parrot')
 app.debug = not helper.is_production() or helper.is_local()
 
@@ -123,6 +122,7 @@ app.register_blueprint(bp_core_dynamic.app)
 app.register_blueprint(bp_core_dynamic_crons.app)
 app.register_blueprint(bp_dashboard.app)
 app.register_blueprint(v1_api.app)
+app.register_blueprint(v1_api_ee.app)
 # Enterprise
 app.register_blueprint(bp_ee.app)
 app.register_blueprint(bp_ee_crons.app)
diff --git a/ee/api/chalicelib/blueprints/app/v1_api_ee.py b/ee/api/chalicelib/blueprints/app/v1_api_ee.py
new file mode 100644
index 000000000..e6f3dc8f2
--- /dev/null
+++ b/ee/api/chalicelib/blueprints/app/v1_api_ee.py
@@ -0,0 +1,14 @@
+from chalice import Blueprint
+
+from chalicelib import _overrides
+from chalicelib.blueprints import bp_authorizers
+from chalicelib.utils import assist_helper
+
+app = Blueprint(__name__)
+_overrides.chalice_app(app)
+
+
+@app.route('/v1/assist/credentials', methods=['GET'], authorizer=bp_authorizers.api_key_authorizer)
+def get_assist_credentials(context):
+    username, credential = assist_helper.get_temporary_credentials()
+    return {"data": {'username': username, 'credential': credential}}
diff --git a/ee/api/chalicelib/blueprints/bp_ee.py b/ee/api/chalicelib/blueprints/bp_ee.py
index 9dae133a0..176e59455 100644
--- a/ee/api/chalicelib/blueprints/bp_ee.py
+++ b/ee/api/chalicelib/blueprints/bp_ee.py
@@ -1,16 +1,9 @@
-import base64
-import hashlib
-import hmac
-from time import time
-
 from chalice import Blueprint
 
 from chalicelib import _overrides
-from chalicelib.blueprints import bp_authorizers
 from chalicelib.core import roles
 from chalicelib.core import unlock
-from chalicelib.utils import helper
-from chalicelib.utils.helper import environ
+from chalicelib.utils import assist_helper
 
 app = Blueprint(__name__)
 _overrides.chalice_app(app)
@@ -60,16 +53,7 @@ def delete_role(roleId, context):
     }
 
 
-@app.route('/v1/assist/credentials', methods=['GET'], authorizer=bp_authorizers.api_key_authorizer)
-@app.route('/assist/credentials', methods=['GET'], authorizer=bp_authorizers.api_key_authorizer)
+@app.route('/assist/credentials', methods=['GET'])
 def get_assist_credentials(context):
-    user = helper.generate_salt()
-    secret = environ["assist_secret"]
-    ttl = int(environ.get("assist_ttl", 48)) * 3600
-    timestamp = int(time()) + ttl
-    username = str(timestamp) + ':' + user
-    dig = hmac.new(bytes(secret, 'utf-8'), bytes(username, 'utf-8'), hashlib.sha1)
-    dig = dig.digest()
-    password = base64.b64encode(dig).decode()
-
-    return {"data": {'username': username, 'password': password}}
+    username, credential = assist_helper.get_temporary_credentials()
+    return {"data": {'username': username, 'credential': credential}}
diff --git a/ee/api/chalicelib/utils/assist_helper.py b/ee/api/chalicelib/utils/assist_helper.py
new file mode 100644
index 000000000..a709fdd2c
--- /dev/null
+++ b/ee/api/chalicelib/utils/assist_helper.py
@@ -0,0 +1,19 @@
+import base64
+import hashlib
+import hmac
+from time import time
+
+from chalicelib.utils import helper
+from chalicelib.utils.helper import environ
+
+
+def get_temporary_credentials():
+    user = helper.generate_salt()
+    secret = environ["assist_secret"]
+    ttl = int(environ.get("assist_ttl", 48)) * 3600
+    timestamp = int(time()) + ttl
+    username = str(timestamp) + ':' + user
+    dig = hmac.new(bytes(secret, 'utf-8'), bytes(username, 'utf-8'), hashlib.sha1)
+    dig = dig.digest()
+    password = base64.b64encode(dig).decode()
+    return user, password