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 os
import sys
from datetime import datetime, timedelta
from pathlib import Path
from typing import Optional
from auth0.authentication.token_verifier import (
@@ -34,14 +36,14 @@ class TokenManager:
self.fernet = Fernet(self.key)
def _get_or_create_key(self):
key_file = "secret.key"
if os.path.exists(key_file):
return open(key_file, "rb").read()
key_filename = "secret.key"
if self.read_secure_file(key_filename):
key = self.read_secure_file(key_filename)
else:
key = Fernet.generate_key()
with open(key_file, "wb") as key_file:
key_file.write(key)
return key
self.save_secure_file(key_filename, key)
return key
def save_tokens(self, access_token, expires_in):
expiration_time = datetime.now() + timedelta(seconds=expires_in)
@@ -50,15 +52,10 @@ class TokenManager:
"expiration": expiration_time.isoformat(),
}
encrypted_data = self.fernet.encrypt(json.dumps(data).encode())
with open(self.file_path, "wb") as file:
file.write(encrypted_data)
self.save_secure_file(self.file_path, encrypted_data)
def get_token(self) -> Optional[str]:
if not os.path.exists(self.file_path):
return None
with open(self.file_path, "rb") as file:
encrypted_data = file.read()
encrypted_data = self.read_secure_file(self.file_path)
decrypted_data = self.fernet.decrypt(encrypted_data)
data = json.loads(decrypted_data)
@@ -68,3 +65,41 @@ class TokenManager:
return None
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:
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]:
return self.storage.search(query)