Update api key file

flojoy/utils.py

@@ -1,6 +1,5 @@
 import decimal
 import json as _json
-
 import numpy as np
 import pandas as pd
 from pathlib import Path
@@ -10,7 +9,7 @@
 import requests
 from dotenv import dotenv_values  # type:ignore
 import difflib
-from typing import Literal
+import keyring
 
 __all__ = [
     "send_to_socket",
@@ -216,81 +215,27 @@ def dump_str(result: Any, limit: int | None = None):
     )
 
 
-def get_frontier_api_key() -> Union[str, None]:
-    home = str(Path.home())
-    api_key = None
-    path = os.path.join(home, ".flojoy/credentials")
-    if not os.path.exists(path):
-        return api_key
-
-    stream = open(path, "r", encoding="utf-8")
-    yaml_dict = yaml.load(stream, Loader=yaml.FullLoader)
-    if yaml_dict is None:
-        return api_key
-    if isinstance(yaml_dict, str) == True:
-        split_by_line = yaml_dict.split("\n")
-        for line in split_by_line:
-            if "FRONTIER_API_KEY" in line:
-                api_key = line.split(":")[1]
+def get_frontier_api_key(name: str, s3Name: str = "") -> Union[dict, None]:
+    if name == "CLOUD":
+        return {"CLOUD": keyring.get_password("system", "CLOUD_API_KEY")}
+    elif name == "OPENAI":
+        return {"OPENAI": keyring.get_password("system", "OPENAI_API_KEY")}
+    elif name == "S3":
+        return {
+            "accessKey": keyring.get_password("system", f"{s3Name}accessKey"),
+            "secretKey": keyring.get_password("system", f"{s3Name}secretKey"),
+        }
     else:
-        api_key = yaml_dict.get("FRONTIER_API_KEY", None)
-    return api_key
-
-
-def set_frontier_api_key(api_key: str):
-    try:
-        home = str(Path.home())
-        file_path = os.path.join(home, ".flojoy/credentials")
-
-        if not os.path.exists(file_path):
-            # Create a new file and write the API_KEY to it
-            with open(file_path, "w") as file:
-                file.write(f"FRONTIER_API_KEY:{api_key}\n")
-        else:
-            # Read the contents of the file
-            with open(file_path, "r") as file:
-                lines = file.readlines()
-
-            # Update the API key if it exists, otherwise append a new line
-            updated = False
-            for i, line in enumerate(lines):
-                if line.startswith("FRONTIER_API_KEY:"):
-                    lines[i] = f"FRONTIER_API_KEY:{api_key}\n"
-                    updated = True
-                    break
+        raise AssertionError("No such API Key exists")
 
-            if not updated:
-                lines.append(f"FRONTIER_API_KEY:{api_key}\n")
-            # Write the updated contents to the file
-            with open(file_path, "w") as file:
-                file.writelines(lines)
 
-    except Exception as e:
-        raise e
-
-
-def set_frontier_s3_key(s3_name: str, s3_access_key: str, s3_secret_key: str):
-    home = str(Path.home())
-    file_path = os.path.join(home, os.path.join(".flojoy", "credentials.yaml"))
-
-    data = {
-        f"{s3_name}": s3_name,
-        f"{s3_name}accessKey": s3_access_key,
-        f"{s3_name}secretKey": s3_secret_key,
-    }
-    if not os.path.exists(file_path):
-        # Create a new file and write the ACCSS_KEY to it
-        with open(file_path, "w") as file:
-            yaml.dump(data, file)
-        return
-
-    # Read the contents of the file
-    with open(file_path, "r") as file:
-        load = yaml.safe_load(file)
-
-    load[f"{s3_name}"] = s3_name
-    load[f"{s3_name}accessKey"] = s3_access_key
-    load[f"{s3_name}secretKey"] = s3_secret_key
-
-    with open(file_path, "w") as file:
-        yaml.dump(load, file)
+def set_frontier_api_key(data: dict, name: str, s3Name: str = ""):
+    if name == "CLOUD":
+        keyring.set_password("system", "CLOUD_API_KEY", data["CLOUD_API_KEY"])
+    elif name == "OPENAI":
+        keyring.set_password("system", "OPENAI_API_KEY", data["OPENAI_API_KEY"])
+    elif name == "S3":
+        keyring.set_password("system", f"{s3Name}accessKey", data[f"{s3Name}accessKey"])
+        keyring.set_password("system", f"{s3Name}secretKey", data[f"{s3Name}secretKey"])
+    else:
+        raise AssertionError("Invalid API Key")