feat: add subsystem saving credentials and remove comment of type hinting

This commit is contained in:
Eduardo Chiarotti
2024-08-22 17:45:27 -03:00
parent 50746d5e8b
commit 156f59760c
2 changed files with 49 additions and 14 deletions

View File

@@ -1,6 +1,8 @@
import json import json
import os import os
import sys
from datetime import datetime, timedelta from datetime import datetime, timedelta
from pathlib import Path
from typing import Optional from typing import Optional
from auth0.authentication.token_verifier import ( from auth0.authentication.token_verifier import (
@@ -34,14 +36,14 @@ class TokenManager:
self.fernet = Fernet(self.key) self.fernet = Fernet(self.key)
def _get_or_create_key(self): def _get_or_create_key(self):
key_file = "secret.key" key_filename = "secret.key"
if os.path.exists(key_file):
return open(key_file, "rb").read() if self.read_secure_file(key_filename):
key = self.read_secure_file(key_filename)
else: else:
key = Fernet.generate_key() key = Fernet.generate_key()
with open(key_file, "wb") as key_file: self.save_secure_file(key_filename, key)
key_file.write(key) return key
return key
def save_tokens(self, access_token, expires_in): def save_tokens(self, access_token, expires_in):
expiration_time = datetime.now() + timedelta(seconds=expires_in) expiration_time = datetime.now() + timedelta(seconds=expires_in)
@@ -50,15 +52,10 @@ class TokenManager:
"expiration": expiration_time.isoformat(), "expiration": expiration_time.isoformat(),
} }
encrypted_data = self.fernet.encrypt(json.dumps(data).encode()) encrypted_data = self.fernet.encrypt(json.dumps(data).encode())
with open(self.file_path, "wb") as file: self.save_secure_file(self.file_path, encrypted_data)
file.write(encrypted_data)
def get_token(self) -> Optional[str]: def get_token(self) -> Optional[str]:
if not os.path.exists(self.file_path): encrypted_data = self.read_secure_file(self.file_path)
return None
with open(self.file_path, "rb") as file:
encrypted_data = file.read()
decrypted_data = self.fernet.decrypt(encrypted_data) decrypted_data = self.fernet.decrypt(encrypted_data)
data = json.loads(decrypted_data) data = json.loads(decrypted_data)
@@ -68,3 +65,41 @@ class TokenManager:
return None return None
return data["access_token"] return data["access_token"]
def get_secure_storage_path(self):
if sys.platform == "win32":
# Windows: Use %LOCALAPPDATA%
base_path = os.environ.get("LOCALAPPDATA")
elif sys.platform == "darwin":
# macOS: Use ~/Library/Application Support
base_path = os.path.expanduser("~/Library/Application Support")
else:
# Linux and other Unix-like: Use ~/.local/share
base_path = os.path.expanduser("~/.local/share")
app_name = "crewai/credentials"
storage_path = Path(base_path) / app_name
storage_path.mkdir(parents=True, exist_ok=True)
return storage_path
def save_secure_file(self, filename, content):
storage_path = self.get_secure_storage_path()
file_path = storage_path / filename
with open(file_path, "wb") as f:
f.write(content)
# Set appropriate permissions (read/write for owner only)
os.chmod(file_path, 0o600)
def read_secure_file(self, filename):
storage_path = self.get_secure_storage_path()
file_path = storage_path / filename
if not file_path.exists():
return None
with open(file_path, "rb") as f:
return f.read()

View File

@@ -21,7 +21,7 @@ class Memory:
if agent: if agent:
metadata["agent"] = agent metadata["agent"] = agent
self.storage.save(value, metadata) # type: ignore # Maybe BUG? Should be self.storage.save(key, value, metadata) self.storage.save(value, metadata)
def search(self, query: str) -> Dict[str, Any]: def search(self, query: str) -> Dict[str, Any]:
return self.storage.search(query) return self.storage.search(query)