token_manager.py

import os
import json
import httpx
from cryptography.fernet import Fernet

class SecureTokenStorage:
    def __init__(self, file_path, key=None):
        self.file_path = file_path
        self.key = key or Fernet.generate_key()
        self.cipher_suite = Fernet(self.key)

    def store_tokens(self, access_token, refresh_token):
        data = {"access_token": access_token, "refresh_token": refresh_token}
        encrypted_data = self.cipher_suite.encrypt(json.dumps(data).encode())

        with open(self.file_path, "wb") as file:
            file.write(encrypted_data)

    def retrieve_tokens(self):
        if not os.path.exists(self.file_path):
            return None, None

        with open(self.file_path, "rb") as file:
            encrypted_data = file.read()

        data = json.loads(self.cipher_suite.decrypt(encrypted_data).decode())
        return data.get("access_token"), data.get("refresh_token")


class TokenManager:
    def __init__(self, client_id, client_secret, port, scopes, token_storage):
        self.client_id = client_id
        self.client_secret = client_secret
        self.port = port
        self.scopes = scopes
        self.token_storage = token_storage

    async def generate_token(self):
        access_token, refresh_token = self.token_storage.retrieve_tokens()

        if not access_token:
            access_token, refresh_token = await self._generate_new_token()

        return access_token, refresh_token

    async def _generate_new_token(self):
        # Implement token generation logic here, similar to your existing generate_token function.
        ...

    async def refresh_tokens(self):
        access_token, refresh_token = await refresh_tokens(
            self.client_id, self.client_secret, self.token_storage.retrieve_refresh_token()
        )

        self.token_storage.store_tokens(access_token, refresh_token)
        return access_token, refresh_token