From f2d1e948d8c981ae6edcd2e4647643ae9474a1a2 Mon Sep 17 00:00:00 2001
From: Daisie Huang <daisieh@gmail.com>
Date: Fri, 22 Dec 2023 18:41:46 -0800
Subject: [PATCH] update to vault store

---
 auth.py              | 11 +++++++++++
 ingest_operations.py |  8 ++++----
 opa_ingest.py        | 32 ++++++++++++++------------------
 3 files changed, 29 insertions(+), 22 deletions(-)

diff --git a/auth.py b/auth.py
index 6fa1ee7..838b8df 100644
--- a/auth.py
+++ b/auth.py
@@ -132,5 +132,16 @@ def is_authed(request: requests.Request):
         return True
     return False
 
+
+def get_opa_access():
+    response, status_code = authx.auth.get_service_store_secret("opa", key="access")
+    return response, status_code
+
+
+def set_opa_access(input):
+    response, status_code = authx.auth.set_service_store_secret("opa", key="access", value=input)
+    return response, status_code
+
+
 if __name__ == "__main__":
     print(get_site_admin_token())
\ No newline at end of file
diff --git a/ingest_operations.py b/ingest_operations.py
index a49a55a..1afcdb0 100644
--- a/ingest_operations.py
+++ b/ingest_operations.py
@@ -56,8 +56,8 @@ def add_s3_credential():
 def add_user_access(program_id, email):
     token = request.headers['Authorization'].split("Bearer ")[1]
     try:
-        result = add_user_to_dataset(email, program_id, token)
-        return result, 200
+        result, status_code = add_user_to_dataset(email, program_id, token)
+        return result, status_code
     except Exception as e:
         return {"error": str(e)}, 500
 
@@ -66,8 +66,8 @@ def add_user_access(program_id, email):
 def remove_user_access(program_id, email):
     token = request.headers['Authorization'].split("Bearer ")[1]
     try:
-        result = remove_user_from_dataset(email, program_id, token)
-        return result, 200
+        result, status_code = remove_user_from_dataset(email, program_id, token)
+        return result, status_code
     except Exception as e:
         return {"error": str(e)}, 500
 
diff --git a/opa_ingest.py b/opa_ingest.py
index 198488c..250e0b6 100644
--- a/opa_ingest.py
+++ b/opa_ingest.py
@@ -12,12 +12,11 @@
 
 
 def add_user_to_dataset(user, dataset, token):
-    headers = {"Authorization": f"Bearer {token}"}
     # get current access:
-    access = requests.get(OPA_URL + "/v1/data/access", headers=headers).json()
-    if "result" not in access:
+    access, status_code = auth.get_opa_access()
+    if status_code != 200:
         raise Exception(f"OPA error: {access}")
-    controlled_access_list = access["result"]["controlled_access_list"]
+    controlled_access_list = access["access"]["controlled_access_list"]
     if user in controlled_access_list:
         if dataset not in controlled_access_list[user]:
             controlled_access_list[user].append(dataset)
@@ -25,29 +24,26 @@ def add_user_to_dataset(user, dataset, token):
         controlled_access_list[user] = [dataset]
 
     # put back:
-    response = requests.put(OPA_URL + "/v1/data/access", headers=headers, json=access["result"])
-    if response.status_code == 204:
-        access = requests.get(OPA_URL + "/v1/data/access", headers=headers).json()
-        return {"access": access["result"]}, 200
-    return {"error": f"{response.status_code}: {response.text}"}, response.status_code
+    response, status_code = auth.set_opa_access(access)
+    if status_code != 200:
+        return {"error": f"{status_code}: {response}"}, status_code
+    return response, 200
 
 
 def remove_user_from_dataset(user, dataset, token):
-    headers = {"Authorization": f"Bearer {token}"}
     # get current access:
-    access = requests.get(OPA_URL + "/v1/data/access", headers=headers).json()
-    if "result" not in access:
+    access, status_code = auth.get_opa_access()
+    if status_code != 200:
         raise Exception(f"OPA error: {access}")
-    controlled_access_list = access["result"]["controlled_access_list"]
+    controlled_access_list = access["access"]["controlled_access_list"]
     if user in controlled_access_list:
         if dataset in controlled_access_list[user]:
             controlled_access_list[user].remove(dataset)
             # put back:
-            response = requests.put(OPA_URL + "/v1/data/access", headers=headers, json=access["result"])
-            if response.status_code == 204:
-                access = requests.get(OPA_URL + "/v1/data/access", headers=headers).json()
-                return {"access": access["result"]}
-            return {"error": f"{response.status_code}: {response.text}"}, response.status_code
+            response, status_code = auth.set_opa_access(access)
+            if status_code != 200:
+                return {"error": f"{status_code}: {response}"}, status_code
+            return access, 200
         return {"error": f"Program {dataset} not authorized for {user}"}, 404
     return {"error": f"User {user} not found"}, 404