From 42d97bef131ecfeceec7c2c6b704cc3d295ad038 Mon Sep 17 00:00:00 2001
From: Taha Yassine Kraiem <tahayk2@gmail.com>
Date: Fri, 18 Nov 2022 13:04:32 +0100
Subject: [PATCH] feat(chalice): vault support old format and new format
 feat(chalice): vault support devtools

---
 api/chalicelib/core/sessions_mobs.py        | 25 +++++------
 ee/api/chalicelib/core/sessions_favorite.py | 49 ++++++++++-----------
 ee/api/chalicelib/utils/s3_extra.py         | 10 +++--
 3 files changed, 41 insertions(+), 43 deletions(-)

diff --git a/api/chalicelib/core/sessions_mobs.py b/api/chalicelib/core/sessions_mobs.py
index 3d966a47c..e2b22dc1b 100644
--- a/api/chalicelib/core/sessions_mobs.py
+++ b/api/chalicelib/core/sessions_mobs.py
@@ -15,6 +15,10 @@ def __get_mob_keys(project_id, session_id):
     ]
 
 
+def __get_mob_keys_deprecated(session_id):
+    return [str(session_id), str(session_id) + "e"]
+
+
 def get_urls(project_id, session_id):
     results = []
     for k in __get_mob_keys(project_id=project_id, session_id=session_id):
@@ -27,23 +31,14 @@ def get_urls(project_id, session_id):
 
 
 def get_urls_depercated(session_id):
-    return [
-        client.generate_presigned_url(
+    results = []
+    for k in __get_mob_keys_deprecated(session_id=session_id):
+        results.append(client.generate_presigned_url(
             'get_object',
-            Params={
-                'Bucket': config("sessions_bucket"),
-                'Key': str(session_id)
-            },
+            Params={'Bucket': config("sessions_bucket"), 'Key': k},
             ExpiresIn=100000
-        ),
-        client.generate_presigned_url(
-            'get_object',
-            Params={
-                'Bucket': config("sessions_bucket"),
-                'Key': str(session_id) + "e"
-            },
-            ExpiresIn=100000
-        )]
+        ))
+    return results
 
 
 def get_ios(session_id):
diff --git a/ee/api/chalicelib/core/sessions_favorite.py b/ee/api/chalicelib/core/sessions_favorite.py
index 3d6496424..7af995bad 100644
--- a/ee/api/chalicelib/core/sessions_favorite.py
+++ b/ee/api/chalicelib/core/sessions_favorite.py
@@ -1,7 +1,7 @@
 from decouple import config
 
 import schemas_ee
-from chalicelib.core import sessions, sessions_favorite_exp
+from chalicelib.core import sessions, sessions_favorite_exp, sessions_mobs, sessions_devtool
 from chalicelib.utils import pg_client, s3_extra
 
 
@@ -34,32 +34,31 @@ def remove_favorite_session(context: schemas_ee.CurrentContext, project_id, sess
 
 
 def favorite_session(context: schemas_ee.CurrentContext, project_id, session_id):
+    keys = sessions_mobs.__get_mob_keys(project_id=project_id, session_id=session_id)
+    keys += sessions_mobs.__get_mob_keys_deprecated(session_id=session_id)  # To support old sessions
+    keys += sessions_devtool.__get_devtools_keys(project_id=project_id, session_id=session_id)
+
     if favorite_session_exists(user_id=context.user_id, session_id=session_id):
-        key = str(session_id)
-        try:
-            s3_extra.tag_file(session_id=key, tag_value=config('RETENTION_D_VALUE', default='default'))
-        except Exception as e:
-            print(f"!!!Error while tagging: {key} to default")
-            print(str(e))
-        key = str(session_id) + "e"
-        try:
-            s3_extra.tag_file(session_id=key, tag_value=config('RETENTION_D_VALUE', default='default'))
-        except Exception as e:
-            print(f"!!!Error while tagging: {key} to default")
-            print(str(e))
+        tag = config('RETENTION_D_VALUE', default='default')
+
+        for k in keys:
+            try:
+                s3_extra.tag_session(file_key=k, tag_value=tag)
+            except Exception as e:
+                print(f"!!!Error while tagging: {k} to {tag} for removal")
+                print(str(e))
+
         return remove_favorite_session(context=context, project_id=project_id, session_id=session_id)
-    key = str(session_id)
-    try:
-        s3_extra.tag_file(session_id=key, tag_value=config('RETENTION_L_VALUE', default='vault'))
-    except Exception as e:
-        print(f"!!!Error while tagging: {key} to vault")
-        print(str(e))
-    key = str(session_id) + "e"
-    try:
-        s3_extra.tag_file(session_id=key, tag_value=config('RETENTION_L_VALUE', default='vault'))
-    except Exception as e:
-        print(f"!!!Error while tagging: {key} to vault")
-        print(str(e))
+
+    tag = config('RETENTION_L_VALUE', default='vault')
+
+    for k in keys:
+        try:
+            s3_extra.tag_session(file_key=k, tag_value=tag)
+        except Exception as e:
+            print(f"!!!Error while tagging: {k} to {tag} for vault")
+            print(str(e))
+
     return add_favorite_session(context=context, project_id=project_id, session_id=session_id)
 
 
diff --git a/ee/api/chalicelib/utils/s3_extra.py b/ee/api/chalicelib/utils/s3_extra.py
index f2a538dcc..4ebf60fec 100644
--- a/ee/api/chalicelib/utils/s3_extra.py
+++ b/ee/api/chalicelib/utils/s3_extra.py
@@ -3,10 +3,14 @@ from decouple import config
 from chalicelib.utils.s3 import client
 
 
-def tag_file(session_id, tag_key='retention', tag_value='vault'):
+def tag_session(file_key, tag_key='retention', tag_value='vault'):
+    return tag_file(file_key=file_key, bucket=config("sessions_bucket"), tag_key=tag_key, tag_value=tag_value)
+
+
+def tag_file(file_key, bucket, tag_key, tag_value):
     return client.put_object_tagging(
-        Bucket=config("sessions_bucket"),
-        Key=session_id,
+        Bucket=bucket,
+        Key=file_key,
         Tagging={
             'TagSet': [
                 {