From 9dec823069d4bcf89289fb4f22adf42dffb3cd34 Mon Sep 17 00:00:00 2001
From: Sepehr Rasouli <sepehrrasouli06@gmail.com>
Date: Tue, 17 Sep 2024 01:57:21 +0330
Subject: [PATCH] added keyring

---
 README.md      | 14 +++-----------
 hey/cli.py     | 11 +++++++++++
 hey/consts.py  |  1 +
 hey/openai.py  | 11 ++++++-----
 pyproject.toml |  3 ++-
 5 files changed, 23 insertions(+), 17 deletions(-)

diff --git a/README.md b/README.md
index afe2707..fcf2613 100644
--- a/README.md
+++ b/README.md
@@ -31,19 +31,11 @@ pip install git+http://github.com/lnxpy/hey.git
 </details>
 
 <details>
-  <summary><h4>2. Set the <code>HEY_TOKEN</code> environment variable</h4></summary>
+  <summary><h4>2. Set the <code>HEY_TOKEN</code></h4></summary>
 
-Once you got the package installed on your system, it's time to add the token that you just copied from [mdb.ai](https://mdb.ai) into either the `.bashrc` (or `.zshrc`) file.
-
-- If you use the default bash shell
-```sh
-echo "export HEY_TOKEN=<TOKEN>" >> ~/.bashrc
-```
-- If you use ZSH
-```sh
-echo "export HEY_TOKEN=<TOKEN>" >> ~/.zshrc
-```
+Once you got the package installed on your system, it's time to add the token that you just copied from [mdb.ai](https://mdb.ai) into hey, use `hey auth` to add your token.
 
+> :warning:: Your token will not be echoed.
 </details>
 
 ### Usage
diff --git a/hey/cli.py b/hey/cli.py
index 9434861..0b56890 100644
--- a/hey/cli.py
+++ b/hey/cli.py
@@ -2,6 +2,8 @@
 from typing import Annotated, Optional
 
 import typer
+import keyring
+import getpass
 from rich.markdown import Markdown
 from rich.panel import Panel
 
@@ -10,6 +12,7 @@
 from hey.consts import APP_NAME
 from hey.editor import open_tmp_editor
 from hey.openai import answer
+from hey.consts import HEY_TOKEN_KEY
 
 app = typer.Typer()
 app.add_typer(cli.app, name="config")
@@ -74,3 +77,11 @@ def ask(
 
     result = answer(user_input, no_style)
     console.print(result)
+
+@app.command()
+def auth():
+    """
+    Take HEY_TOKEN from user.
+    """
+    password = getpass.getpass("Copy and paste your token here, [Token will not be echoed] > ")
+    keyring.set_password("system",HEY_TOKEN_KEY,password)
diff --git a/hey/consts.py b/hey/consts.py
index 76771d3..dd9b281 100644
--- a/hey/consts.py
+++ b/hey/consts.py
@@ -5,6 +5,7 @@
 
 # App name
 APP_NAME = "Hey"
+HEY_TOKEN_KEY = "HEY_TOKEN"
 
 if platform.system().lower() == "windows":
     DEFAULT_EDITOR = os.environ.get("EDITOR", "notepad")
diff --git a/hey/openai.py b/hey/openai.py
index e832809..e0948ee 100644
--- a/hey/openai.py
+++ b/hey/openai.py
@@ -1,5 +1,6 @@
 import os
 import time
+import keyring
 
 from openai import OpenAI, OpenAIError
 from rich.console import Console
@@ -8,6 +9,7 @@
 
 from hey.configs import configs
 from hey.consts import BASE_CONFIG
+from hey.consts import HEY_TOKEN_KEY
 from hey.exceptions import ConnectionIssue, TokenIsNotDefined
 
 
@@ -16,14 +18,13 @@ def __init__(self) -> None:
         self.is_valid = False
 
     def validate(self) -> str:
-        token = os.environ.get("HEY_TOKEN", None)
+        token = keyring.get_password("system",HEY_TOKEN_KEY)
         if token:
             self.is_valid = True
             return token
-        else:
-            raise TokenIsNotDefined(
-                "make sure the `HEY_TOKEN` is defined in the .bashrc/.zshrc file."
-            )
+        raise TokenIsNotDefined(
+            "token is not defined, Use `hey auth` to set your token."
+        )
 
 
 class ChatGPT:
diff --git a/pyproject.toml b/pyproject.toml
index b247959..d947b82 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -12,7 +12,8 @@ dependencies = [
     "typer >= 0.12.3",
     "rich >= 13.7.1",
     "openai >= 1.30.1",
-    "platformdirs >= 4.2.2"
+    "platformdirs >= 4.2.2",
+    "keyring >= 25.3.0"
 ]
 readme = { file = "README.md", content-type = "text/markdown" }
 license = { file = "LICENSE" }