neuronum 8.4.0__py3-none-any.whl → 10.0.0__py3-none-any.whl

neuronum/neuronum.py

@@ -1,338 +1,745 @@
 import aiohttp
-from typing import AsyncGenerator
+import aiofiles
+from typing import AsyncGenerator, Optional, Dict, Any, List
 import websockets
 import json
 import asyncio
 import base64
+import ssl
 import os
+import time
+import logging
 from pathlib import Path
-from websockets.exceptions import ConnectionClosed
+from dataclasses import dataclass
+from websockets.exceptions import ConnectionClosed, WebSocketException
 from cryptography.hazmat.primitives.asymmetric import ec
 from cryptography.hazmat.primitives import serialization, hashes
 from cryptography.hazmat.primitives.ciphers.aead import AESGCM
 from cryptography.hazmat.primitives.kdf.hkdf import HKDF
 from cryptography.hazmat.backends import default_backend
+from abc import ABC, abstractmethod
 
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
 
-class Node:
-    def __init__(self, id: str, private_key: str, public_key: str):
-        self.node_id = id
-        self.private_key_path = private_key
-        self.public_key_path = public_key
-        self.queue = asyncio.Queue()
-        self.host = self._load_host()
-        self.network = self._load_network()
-        self.synapse = self._load_synapse()
-        self.password = self._load_password()
-        self._private_key = self._load_private_key()
-        self._public_key = self._load_public_key()
 
-
-    def to_dict(self) -> dict:
-        return {
-            "host": self.host,
-            "password": self.password,
-            "synapse": self.synapse
-        }
+# Custom Exceptions
+class NeuronumError(Exception):
+    """Base exception for Neuronum errors"""
+    pass
 
 
-    def _load_private_key(self):
-        try:
-            with open(self.private_key_path, "rb") as f:
-                private_key = serialization.load_pem_private_key(
-                    f.read(),
-                    password=None,
-                    backend=default_backend()
-                )
-                return private_key
-        except FileNotFoundError:
-            print(f"Private key file not found at {self.private_key_path}.")
-            return None
-        
+class AuthenticationError(NeuronumError):
+    """Raised when authentication fails"""
+    pass
 
-    def _load_host(self):
-        credentials_folder_path = Path.home() / ".neuronum"
-        env_path = credentials_folder_path / ".env"
 
-        env_data = {}  
+class EncryptionError(NeuronumError):
+    """Raised when encryption/decryption fails"""
+    pass
 
-        try:
-            with open(env_path, "r") as f:
-                for line in f:
-                    key, value = line.strip().split("=")
-                    env_data[key] = value
 
-            host = env_data.get("HOST", "")
-            return host
-        except FileNotFoundError:
-            print(f"Cell Host not found")
-            return None
-        
+class CellNotFoundError(NeuronumError):
+    """Raised when a cell cannot be found"""
+    pass
 
-    def _load_password(self):
-        credentials_folder_path = Path.home() / ".neuronum"
-        env_path = credentials_folder_path / ".env"
 
-        env_data = {}  
+class NetworkError(NeuronumError):
+    """Raised when network operations fail"""
+    pass
 
-        try:
-            with open(env_path, "r") as f:
-                for line in f:
-                    key, value = line.strip().split("=")
-                    env_data[key] = value
 
-            host = env_data.get("PASSWORD", "")
-            return host
-        except FileNotFoundError:
-            print(f"Cell Password not found")
-            return None
-        
-    def _load_synapse(self):
-        credentials_folder_path = Path.home() / ".neuronum"
-        env_path = credentials_folder_path / ".env"
+# Configuration
+@dataclass
+class ClientConfig:
+    """Client configuration settings"""
+    network: str = "neuronum.net"
+    cache_expiry: int = 3600
+    credentials_path: Path = Path.home() / ".neuronum"
+    timeout: int = 30
+    max_retries: int = 3
+    retry_delay: float = 1.0
+    max_retry_delay: float = 60.0
 
-        env_data = {}  
 
+class CryptoManager:
+    """Handles all cryptographic operations"""
+    
+    def __init__(self, private_key: Optional[ec.EllipticCurvePrivateKey] = None):
+        self._private_key = private_key
+        self._public_key = private_key.public_key() if private_key else None
+    
+    def sign_message(self, message: bytes) -> str:
+        """Sign a message with the private key"""
+        if not self._private_key:
+            raise EncryptionError("Private key not available for signing")
+        
         try:
-            with open(env_path, "r") as f:
-                for line in f:
-                    key, value = line.strip().split("=")
-                    env_data[key] = value
-
-            host = env_data.get("SYNAPSE", "")
-            return host
-        except FileNotFoundError:
-            print(f"Cell Synapse not found")
-            return None
+            signature = self._private_key.sign(message, ec.ECDSA(hashes.SHA256()))
+            return base64.b64encode(signature).decode()
+        except Exception as e:
+            logger.error("Failed to sign message", exc_info=True)
+            raise EncryptionError(f"Message signing failed: {e}")
+    
+    def get_public_key_pem(self) -> str:
+        """Get public key in PEM format"""
+        if not self._public_key:
+            raise EncryptionError("Public key not available")
         
-    def _load_network(self):
-        credentials_folder_path = Path.home() / ".neuronum"
-        env_path = credentials_folder_path / ".env"
-
-        env_data = {}  
-
+        pem_bytes = self._public_key.public_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PublicFormat.SubjectPublicKeyInfo
+        )
+        return pem_bytes.decode('utf-8')
+    
+    def load_public_key_from_pem(self, pem_string: str) -> ec.EllipticCurvePublicKey:
+        """Load a public key from PEM format"""
         try:
-            with open(env_path, "r") as f:
-                for line in f:
-                    key, value = line.strip().split("=")
-                    env_data[key] = value
-
-            host = env_data.get("NETWORK", "")
-            return host
-        except FileNotFoundError:
-            print(f"Cell Network not found")
-            return None
-
-
-    def _load_public_key(self):
+            return serialization.load_pem_public_key(
+                pem_string.encode(), 
+                backend=default_backend()
+            )
+        except Exception as e:
+            logger.error("Failed to load public key from PEM", exc_info=True)
+            raise EncryptionError(f"Failed to load public key: {e}")
+    
+    @staticmethod
+    def safe_b64decode(data: str) -> bytes:
+        """Safely decode base64 with proper padding"""
+        padding = 4 - (len(data) % 4)
+        if padding != 4:
+            data += '=' * padding
+        return base64.urlsafe_b64decode(data)
+    
+    def encrypt_with_ecdh_aesgcm(
+        self, 
+        public_key: ec.EllipticCurvePublicKey, 
+        plaintext_dict: Dict[str, Any]
+    ) -> Dict[str, str]:
+        """Encrypt data using ECDH + AES-GCM"""
         try:
-            with open(self.public_key_path, "rb") as f:
-                public_key = serialization.load_pem_public_key(
-                    f.read(),
-                    backend=default_backend()
-                )
-                return public_key
-        except FileNotFoundError:
-            print(f"Public key file not found at {self.public_key_path}. Deriving from private key.")
-            if self._private_key:
-                return self._private_key.public_key()
-            else:
-                return None
-
-
-    def get_public_key_jwk(self):
-        public_key = self._load_public_key()
-        if not public_key:
-            print("Public key not loaded. Cannot generate JWK.")
-            return None
-        
-        public_numbers = public_key.public_numbers()
-        
-        x_bytes = public_numbers.x.to_bytes((public_numbers.x.bit_length() + 7) // 8, 'big')
-        y_bytes = public_numbers.y.to_bytes((public_numbers.y.bit_length() + 7) // 8, 'big')
-        
-        return {
-            "kty": "EC",
-            "crv": "P-256",
-            "x": base64.urlsafe_b64encode(x_bytes).rstrip(b'=').decode('utf-8'),
-            "y": base64.urlsafe_b64encode(y_bytes).rstrip(b'=').decode('utf-8')
-        }
+            ephemeral_private = ec.generate_private_key(ec.SECP256R1())
+            shared_secret = ephemeral_private.exchange(ec.ECDH(), public_key)
+            derived_key = HKDF(
+                algorithm=hashes.SHA256(), 
+                length=32, 
+                salt=None, 
+                info=b'handshake data'
+            ).derive(shared_secret)
+            
+            aesgcm = AESGCM(derived_key)
+            nonce = os.urandom(12)
+            plaintext_bytes = json.dumps(plaintext_dict).encode()
+            ciphertext = aesgcm.encrypt(nonce, plaintext_bytes, None)
+            
+            ephemeral_public_bytes = ephemeral_private.public_key().public_bytes(
+                serialization.Encoding.X962, 
+                serialization.PublicFormat.UncompressedPoint
+            )
+            
+            return {
+                'ciphertext': base64.urlsafe_b64encode(ciphertext).rstrip(b'=').decode(),
+                'nonce': base64.urlsafe_b64encode(nonce).rstrip(b'=').decode(),
+                'ephemeralPublicKey': base64.urlsafe_b64encode(ephemeral_public_bytes).rstrip(b'=').decode()
+            }
+        except Exception as e:
+            logger.error("Encryption failed", exc_info=True)
+            raise EncryptionError(f"Encryption failed: {e}")
     
-  
-    def _decrypt_with_ecdh_aesgcm(self, ephemeral_public_key_bytes, nonce, ciphertext):
+    def decrypt_with_ecdh_aesgcm(
+        self, 
+        ephemeral_public_key_bytes: bytes, 
+        nonce: bytes, 
+        ciphertext: bytes
+    ) -> Dict[str, Any]:
+        """Decrypt data using ECDH + AES-GCM"""
+        if not self._private_key:
+            raise EncryptionError("Private key not available for decryption")
+        
         try:
             ephemeral_public_key = ec.EllipticCurvePublicKey.from_encoded_point(
                 ec.SECP256R1(), ephemeral_public_key_bytes
             )
-
             shared_secret = self._private_key.exchange(ec.ECDH(), ephemeral_public_key)
-            
             derived_key = HKDF(
-                algorithm=hashes.SHA256(),
-                length=32,
-                salt=None,
+                algorithm=hashes.SHA256(), 
+                length=32, 
+                salt=None, 
                 info=b'handshake data'
             ).derive(shared_secret)
-
+            
             aesgcm = AESGCM(derived_key)
             plaintext_bytes = aesgcm.decrypt(nonce, ciphertext, None)
             return json.loads(plaintext_bytes.decode())
-            
         except Exception as e:
-            print(f"Decryption failed: {e}")
+            logger.error("Decryption failed", exc_info=True)
+            raise EncryptionError(f"Decryption failed: {e}")
+
+
+class CacheManager:
+    """Manages cell cache with async file operations"""
+    
+    def __init__(self, config: ClientConfig):
+        self.config = config
+        self.cache_file = config.credentials_path / "cells.json"
+        self._lock = asyncio.Lock()
+        self._memory_cache: Optional[List[Dict[str, Any]]] = None
+        self._cache_time: Optional[float] = None
+    
+    async def get_cells(self) -> List[Dict[str, Any]]:
+        """Get cached cells if valid, otherwise fetch new"""
+        async with self._lock:
+            # Check memory cache first
+            if self._is_memory_cache_valid():
+                logger.debug("Using in-memory cache")
+                return self._memory_cache
+            
+            # Check file cache
+            if await self._is_file_cache_valid():
+                logger.debug("Using file cache")
+                cells = await self._load_from_file()
+                self._update_memory_cache(cells)
+                return cells
+            
             return None
+    
+    async def update_cells(self, cells: List[Dict[str, Any]]) -> None:
+        """Update cache with new cell data"""
+        async with self._lock:
+            self._update_memory_cache(cells)
+            await self._save_to_file(cells)
+    
+    def _is_memory_cache_valid(self) -> bool:
+        """Check if memory cache is still valid"""
+        if not self._memory_cache or not self._cache_time:
+            return False
+        return (time.time() - self._cache_time) < self.config.cache_expiry
+    
+    async def _is_file_cache_valid(self) -> bool:
+        """Check if file cache is still valid"""
+        if not self.cache_file.exists():
+            return False
+        
+        try:
+            file_mtime = os.path.getmtime(self.cache_file)
+            return (time.time() - file_mtime) < self.config.cache_expiry
+        except OSError:
+            return False
+    
+    async def _load_from_file(self) -> List[Dict[str, Any]]:
+        """Load cells from cache file"""
+        try:
+            async with aiofiles.open(self.cache_file, 'r') as f:
+                content = await f.read()
+                return json.loads(content)
+        except (FileNotFoundError, json.JSONDecodeError) as e:
+            logger.warning(f"Failed to load cache file: {e}")
+            return []
+    
+    async def _save_to_file(self, cells: List[Dict[str, Any]]) -> None:
+        """Save cells to cache file"""
+        try:
+            self.config.credentials_path.mkdir(parents=True, exist_ok=True)
+            async with aiofiles.open(self.cache_file, 'w') as f:
+                await f.write(json.dumps(cells, indent=4))
+            logger.debug("Cache file updated")
+        except Exception as e:
+            logger.error(f"Failed to save cache file: {e}")
+    
+    def _update_memory_cache(self, cells: List[Dict[str, Any]]) -> None:
+        """Update in-memory cache"""
+        self._memory_cache = cells
+        self._cache_time = time.time()
 
 
-    async def sync(self) -> AsyncGenerator[str, None]:
-        full_url = f"wss://{self.network}/sync/{self.node_id}"
-        auth_payload = {
+class NetworkClient:
+    """Handles all network operations with retry logic"""
+    
+    def __init__(self, config: ClientConfig):
+        self.config = config
+        self._session: Optional[aiohttp.ClientSession] = None
+    
+    async def __aenter__(self):
+        self._session = aiohttp.ClientSession(
+            timeout=aiohttp.ClientTimeout(total=self.config.timeout)
+        )
+        return self
+    
+    async def __aexit__(self, *args):
+        if self._session:
+            await self._session.close()
+    
+    async def post_request(
+        self, 
+        url: str, 
+        payload: Dict[str, Any],
+        retry_count: int = 0
+    ) -> Optional[Dict[str, Any]]:
+        """Make POST request with retry logic"""
+        if not self._session:
+            self._session = aiohttp.ClientSession(
+                timeout=aiohttp.ClientTimeout(total=self.config.timeout)
+            )
+        
+        try:
+            async with self._session.post(url, json=payload) as response:
+                response.raise_for_status()
+                return await response.json()
+        except aiohttp.ClientResponseError as e:
+            logger.error(f"HTTP error {e.status} for URL: {url}")
+            if retry_count < self.config.max_retries and e.status >= 500:
+                return await self._retry_request(url, payload, retry_count)
+            raise NetworkError(f"HTTP {e.status} error")
+        except aiohttp.ClientError as e:
+            logger.error(f"Client error for URL {url}: {e}")
+            if retry_count < self.config.max_retries:
+                return await self._retry_request(url, payload, retry_count)
+            raise NetworkError(f"Client error: {e}")
+        except Exception as e:
+            logger.error(f"Unexpected error for URL {url}: {e}")
+            raise NetworkError(f"Unexpected error: {e}")
+    
+    async def _retry_request(
+        self, 
+        url: str, 
+        payload: Dict[str, Any], 
+        retry_count: int
+    ) -> Optional[Dict[str, Any]]:
+        """Retry request with exponential backoff"""
+        delay = min(
+            self.config.retry_delay * (2 ** retry_count),
+            self.config.max_retry_delay
+        )
+        logger.info(f"Retrying request in {delay}s (attempt {retry_count + 1})")
+        await asyncio.sleep(delay)
+        return await self.post_request(url, payload, retry_count + 1)
+
+
+class BaseClient(ABC):
+    """Base client with common functionality"""
+    
+    def __init__(self, config: Optional[ClientConfig] = None):
+        self.config = config or ClientConfig()
+        self.env: Dict[str, str] = {}
+        self._crypto: Optional[CryptoManager] = None
+        self._cache_manager = CacheManager(self.config)
+        self._network_client = NetworkClient(self.config)
+        self.host = ""
+        self.network = self.config.network
+    
+    @abstractmethod
+    def _load_private_key(self) -> Optional[ec.EllipticCurvePrivateKey]:
+        """Load private key - to be implemented by subclasses"""
+        pass
+    
+    def _init_crypto(self, private_key: Optional[ec.EllipticCurvePrivateKey]) -> None:
+        """Initialize crypto manager with private key"""
+        self._crypto = CryptoManager(private_key)
+    
+    def to_dict(self) -> Dict[str, str]:
+        """Create authentication payload"""
+        if not self._crypto:
+            logger.warning("Crypto manager not initialized")
+            timestamp = str(int(time.time()))
+            return {
+                "host": self.host,
+                "signed_message": "",
+                "message": f"host={self.host};timestamp={timestamp}"
+            }
+        
+        timestamp = str(int(time.time()))
+        message = f"host={self.host};timestamp={timestamp}"
+        
+        try:
+            signed_message = self._crypto.sign_message(message.encode())
+        except EncryptionError:
+            logger.error("Failed to sign authentication message")
+            signed_message = ""
+        
+        return {
             "host": self.host,
-            "password": self.password,
-            "synapse": self.synapse,
+            "signed_message": signed_message,
+            "message": message
         }
+    
+    async def _get_target_cell_public_key(self, cell_id: str) -> str:
+        """Get public key for target cell"""
+        # Try cached cells first
+        cells = await self.list_cells(update=False)
+        
+        for cell in cells:
+            if cell.get('cell_id') == cell_id:
+                public_key = cell.get('public_key', {})
+                if public_key:
+                    return public_key
+        
+        # Refresh cache and try again
+        logger.info(f"Cell {cell_id} not in cache, refreshing")
+        cells = await self.list_cells(update=True)
+        
+        for cell in cells:
+            if cell.get('cell_id') == cell_id:
+                public_key = cell.get('public_key', {})
+                if public_key:
+                    return public_key
+        
+        raise CellNotFoundError(f"Cell not found: {cell_id}")
+    
+    async def list_cells(self, update: bool = False) -> List[Dict[str, Any]]:
+        """
+        List all available cells.
+        
+        Args:
+            update: Force cache refresh if True
+            
+        Returns:
+            List of cell information dictionaries
+        """
+        if not update:
+            cached_cells = await self._cache_manager.get_cells()
+            if cached_cells is not None:
+                return cached_cells
+        
+        # Fetch from API
+        full_url = f"https://{self.network}/api/list_cells"
+        payload = {"cell": self.to_dict()}
+        
+        try:
+            data = await self._network_client.post_request(full_url, payload)
+            cells = data.get("Cells", []) if data else []
+            await self._cache_manager.update_cells(cells)
+            return cells
+        except NetworkError as e:
+            logger.error(f"Failed to fetch cells: {e}")
+            return []
+    
+    async def tx_response(
+        self, 
+        transmitter_id: str, 
+        data: Dict[str, Any], 
+        client_public_key_str: str
+    ) -> None:
+        """
+        Send an encrypted response to a transmitter.
+        
+        Args:
+            transmitter_id: ID of the transmitter
+            data: Response data to encrypt
+            client_public_key_str: Client's public key in PEM format
+            
+        Raises:
+            EncryptionError: If encryption fails
+            NetworkError: If network request fails
+        """
+        if not self._crypto:
+            raise EncryptionError("Crypto manager not initialized")
+        
+        if not client_public_key_str:
+            raise ValueError("client_public_key_str is required")
+        
+        url = f"https://{self.network}/api/tx_response/{transmitter_id}"
+        
+        public_key = self._crypto.load_public_key_from_pem(client_public_key_str)
+        encrypted_payload = self._crypto.encrypt_with_ecdh_aesgcm(public_key, data)
+        payload = {"data": encrypted_payload, "cell": self.to_dict()}
+        
+        await self._network_client.post_request(url, payload)
+        logger.info(f"Response sent to transmitter {transmitter_id}")
+    
+    async def activate_tx(
+        self, 
+        cell_id: str, 
+        data: Dict[str, Any]
+    ) -> Optional[Dict[str, Any]]:
+        """
+        Activate a transaction with the specified cell.
+        
+        Args:
+            cell_id: Target cell identifier
+            data: Transaction data to encrypt and send
+            
+        Returns:
+            Decrypted response dict or None on failure
+            
+        Raises:
+            CellNotFoundError: If cell_id is not found
+            EncryptionError: If encryption/decryption fails
+            NetworkError: If network request fails
+        """
+        if not self._crypto:
+            raise EncryptionError("Crypto manager not initialized")
+        
+        url = f"https://{self.network}/api/activate_tx/{cell_id}"
+        payload = {"cell": self.to_dict()}
+        
+        # Get target cell's public key
+        public_key_pem_str = await self._get_target_cell_public_key(cell_id)
+        public_key_object = self._crypto.load_public_key_from_pem(public_key_pem_str)
+        
+        # Prepare and encrypt data
+        data_to_encrypt = data.copy()
+        data_to_encrypt["public_key"] = self._crypto.get_public_key_pem()
+        encrypted_payload = self._crypto.encrypt_with_ecdh_aesgcm(
+            public_key_object, 
+            data_to_encrypt
+        )
+        payload["data"] = {"encrypted": encrypted_payload}
+        
+        response_data = await self._network_client.post_request(url, payload)
+        
+        if not response_data or "response" not in response_data:
+            logger.warning("Unexpected or missing response")
+            return response_data
+        
+        inner_response = response_data["response"]
+        
+        # Decrypt response if encrypted
+        if "ciphertext" in inner_response:
+            try:
+                ephemeral_public_key_bytes = CryptoManager.safe_b64decode(
+                    inner_response["ephemeralPublicKey"]
+                )
+                nonce = CryptoManager.safe_b64decode(inner_response["nonce"])
+                ciphertext = CryptoManager.safe_b64decode(inner_response["ciphertext"])
+                
+                return self._crypto.decrypt_with_ecdh_aesgcm(
+                    ephemeral_public_key_bytes, nonce, ciphertext
+                )
+            except EncryptionError:
+                logger.error("Failed to decrypt response")
+                return None
+        else:
+            logger.debug("Received unencrypted response")
+            return inner_response
+    
+    async def sync(self) -> AsyncGenerator[Dict[str, Any], None]:
+        """
+        Sync with the network and yield operations as they arrive.
+        
+        Yields:
+            Operation dictionaries with decrypted data
+            
+        Raises:
+            ValueError: If not called from Cell instance or host not set
+        """
+        if not isinstance(self, Cell):
+            raise ValueError("sync must be called from a Cell instance")
+        
+        cell = getattr(self, 'host', None)
+        if not cell:
+            raise ValueError("host is required for Cell sync")
+        
+        full_url = f"wss://{self.network}/sync/{cell}"
+        auth_payload = self.to_dict()
+        
+        logger.info(f"Starting sync with {cell}")
+        
+        retry_count = 0
         while True:
             try:
-                async with websockets.connect(full_url) as ws:
+                ssl_context = ssl.create_default_context()
+                
+                async with websockets.connect(full_url, ssl=ssl_context) as ws:
                     await ws.send(json.dumps(auth_payload))
-                    print("Node syncing...")
+                    logger.info(f"Connected and authenticated to {cell}")
+                    retry_count = 0  # Reset on successful connection
+                    
                     while True:
                         try:
-                            raw_operation = await ws.recv()
+                            raw_operation = await asyncio.wait_for(
+                                ws.recv(), 
+                                timeout=self.config.timeout
+                            )
                             operation = json.loads(raw_operation)
                             
                             if "encrypted" in operation.get("data", {}):
                                 encrypted_data = operation["data"]["encrypted"]
                                 
-                                ephemeral_public_key_b64 = encrypted_data["ephemeralPublicKey"]
-                                ephemeral_public_key_b64 += '=' * ((4 - len(ephemeral_public_key_b64) % 4) % 4)
-                                ephemeral_public_key_bytes = base64.urlsafe_b64decode(ephemeral_public_key_b64)
-
-                                nonce_b64 = encrypted_data["nonce"]
-                                nonce_b64 += '=' * ((4 - len(nonce_b64) % 4) % 4)
-                                nonce = base64.urlsafe_b64decode(nonce_b64)
-                                
-                                ciphertext_b64 = encrypted_data["ciphertext"]
-                                ciphertext_b64 += '=' * ((4 - len(ciphertext_b64) % 4) % 4)
-                                ciphertext = base64.urlsafe_b64decode(ciphertext_b64)
-                                
-                                decrypted_data = self._decrypt_with_ecdh_aesgcm(
-                                    ephemeral_public_key_bytes, nonce, ciphertext
-                                )
-                                
-                                if decrypted_data:
+                                try:
+                                    ephemeral_public_key_bytes = CryptoManager.safe_b64decode(
+                                        encrypted_data["ephemeralPublicKey"]
+                                    )
+                                    nonce = CryptoManager.safe_b64decode(
+                                        encrypted_data["nonce"]
+                                    )
+                                    ciphertext = CryptoManager.safe_b64decode(
+                                        encrypted_data["ciphertext"]
+                                    )
+                                    
+                                    decrypted_data = self._crypto.decrypt_with_ecdh_aesgcm(
+                                        ephemeral_public_key_bytes, nonce, ciphertext
+                                    )
+                                    
                                     operation["data"].update(decrypted_data)
                                     operation["data"].pop("encrypted")
                                     yield operation
-                                else:
-                                    print("Failed to decrypt incoming data. Skipping...")
+                                except EncryptionError:
+                                    logger.error("Failed to decrypt operation")
                             else:
-                                yield operation
+                                logger.warning("Received unencrypted data")
+                        
                         except asyncio.TimeoutError:
-                            print("No data received. Continuing...")
+                            await ws.ping()
+                            continue
                         except ConnectionClosed as e:
-                            if e.code == 1000:
-                                print(f"WebSocket closed cleanly (code 1000). Reconnecting...")
-                            else:
-                                print(f"Connection closed with error code {e.code}: {e.reason}. Reconnecting...")
+                            logger.warning(f"Connection closed: {e.code} - {e.reason}")
                             break
                         except Exception as e:
-                            print(f"Unexpected error in recv loop: {e}")
+                            logger.error(f"Error in receive loop: {e}")
                             break
-            except websockets.exceptions.WebSocketException as e:
-                print(f"WebSocket error occurred: {e}. Retrying in 5 seconds...")
+            
+            except WebSocketException as e:
+                logger.error(f"WebSocket error: {e}")
             except Exception as e:
-                print(f"General error occurred: {e}. Retrying in 5 seconds...")
-            await asyncio.sleep(3)
-
+                logger.error(f"General error in sync: {e}")
+            
+            # Exponential backoff for reconnection
+            retry_count += 1
+            delay = min(
+                self.config.retry_delay * (2 ** retry_count),
+                self.config.max_retry_delay
+            )
+            logger.info(f"Reconnecting in {delay}s")
+            await asyncio.sleep(delay)
+    
+    async def stream(self, cell_id: str, data: Dict[str, Any]) -> bool:
+        """
+        Stream data to a target cell.
+        
+        Args:
+            cell_id: Target cell identifier
+            data: Data to encrypt and stream
+            
+        Returns:
+            True if successful, False otherwise
+            
+        Raises:
+            ValueError: If not called from Cell instance or host not set
+            CellNotFoundError: If cell_id is not found
+            EncryptionError: If encryption fails
+        """
+        if not isinstance(self, Cell):
+            raise ValueError("stream must be called from a Cell instance")
+        
+        if not getattr(self, 'host', None):
+            raise ValueError("host is required for Cell stream")
+        
+        if not self._crypto:
+            raise EncryptionError("Crypto manager not initialized")
+        
+        # Get target cell's public key
+        public_key_pem_str = await self._get_target_cell_public_key(cell_id)
+        public_key_object = self._crypto.load_public_key_from_pem(public_key_pem_str)
+        
+        # Prepare and encrypt data
+        data_to_encrypt = data.copy()
+        data_to_encrypt["public_key"] = self._crypto.get_public_key_pem()
+        encrypted_payload = self._crypto.encrypt_with_ecdh_aesgcm(
+            public_key_object, 
+            data_to_encrypt
+        )
+        
+        auth_payload = self.to_dict()
+        data_payload = {"data": {"encrypted": encrypted_payload}}
+        send_payload = {**auth_payload, **data_payload}
+        
+        full_url = f"wss://{self.network}/stream/{cell_id}"
+        
+        try:
+            ssl_context = ssl.create_default_context()
+            async with websockets.connect(full_url, ssl=ssl_context) as ws:
+                await ws.send(json.dumps(send_payload))
+                logger.info(f"Data streamed to {cell_id}")
+                
+                try:
+                    ack = await asyncio.wait_for(ws.recv(), timeout=2)
+                    logger.debug(f"Server acknowledgment: {ack}")
+                except asyncio.TimeoutError:
+                    logger.debug("No immediate acknowledgment (data sent)")
+                except Exception as e:
+                    logger.warning(f"Error reading acknowledgment: {e}")
+                
+                return True
+        except WebSocketException as e:
+            logger.error(f"WebSocket error during stream: {e}")
+            return False
+        except Exception as e:
+            logger.error(f"Error during stream: {e}")
+            return False
 
-    async def list_nodes(self):
-        full_url = f"https://{self.network}/api/list_nodes"
-        list_nodes_payload = {
-            "cell": self.to_dict()
-        }
-        async with aiohttp.ClientSession() as session:
-            try:
-                async with session.post(full_url, json=list_nodes_payload) as response:
-                    response.raise_for_status()
-                    data = await response.json()
-                    return data.get("Nodes", [])
-            except aiohttp.ClientError as e:
-                print(f"Error sending request: {e}")
-            except Exception as e:
-                print(f"Unexpected error: {e}")
-  
 
-    def _load_public_key_from_jwk(self, jwk):
+class Cell(BaseClient):
+    """Cell client implementation"""
+    
+    def __init__(self, config: Optional[ClientConfig] = None):
+        super().__init__(config)
+        self.env = self._load_env()
+        private_key = self._load_private_key()
+        self._init_crypto(private_key)
+        
+        self.host = self.env.get("HOST", "")
+        if not self.host:
+            logger.warning("HOST not set in environment")
+    
+    def _load_private_key(self) -> Optional[ec.EllipticCurvePrivateKey]:
+        """Load private key from credentials folder"""
+        credentials_path = self.config.credentials_path
+        credentials_path.mkdir(parents=True, exist_ok=True)
+        
+        key_path = credentials_path / "private_key.pem"
+        
         try:
-            print(jwk)
-            x = base64.urlsafe_b64decode(jwk['x'] + '==')
-            y = base64.urlsafe_b64decode(jwk['y'] + '==')
-            public_numbers = ec.EllipticCurvePublicNumbers(
-                int.from_bytes(x, 'big'),
-                int.from_bytes(y, 'big'),
-                ec.SECP256R1()
-            )
-            return public_numbers.public_key(default_backend())
-        except (KeyError, ValueError, TypeError) as e:
-            print(f"Error loading public key from JWK string: {e}")
+            with open(key_path, "rb") as f:
+                private_key = serialization.load_pem_private_key(
+                    f.read(), 
+                    password=None, 
+                    backend=default_backend()
+                )
+            
+            # Check file permissions (should be 0600)
+            stat = os.stat(key_path)
+            if stat.st_mode & 0o177:
+                logger.warning(
+                    f"Private key file has insecure permissions: {oct(stat.st_mode)}"
+                )
+            
+            logger.info("Private key loaded successfully")
+            return private_key
+        except FileNotFoundError:
+            logger.error(f"Private key not found at {key_path}")
+            return None
+        except Exception as e:
+            logger.error(f"Error loading private key: {e}")
             return None
-
-
-    def _encrypt_with_ecdh_aesgcm(self, public_key, plaintext_dict):
-        ephemeral_private = ec.generate_private_key(ec.SECP256R1())
-        ephemeral_public = ephemeral_private.public_key()
-        shared_secret = ephemeral_private.exchange(ec.ECDH(), public_key)
-        derived_key = HKDF(
-            algorithm=hashes.SHA256(),
-            length=32,
-            salt=None,
-            info=b'handshake data'
-        ).derive(shared_secret)
-        aesgcm = AESGCM(derived_key)
-        nonce = os.urandom(12)
-        plaintext_bytes = json.dumps(plaintext_dict).encode()
-        ciphertext = aesgcm.encrypt(nonce, plaintext_bytes, None)
-        ephemeral_public_bytes = ephemeral_public.public_bytes(
-            serialization.Encoding.X962,
-            serialization.PublicFormat.UncompressedPoint
-        )
-        return {
-            'ciphertext': base64.b64encode(ciphertext).decode(),
-            'nonce': base64.b64encode(nonce).decode(),
-            'ephemeralPublicKey': base64.b64encode(ephemeral_public_bytes).decode()
-        }
     
-
-    async def tx_response(self, transmitter_id: str, data: dict, client_public_key_str):
-        if isinstance(client_public_key_str, str):
-            try:
-                client_public_key_jwk = json.loads(client_public_key_str)
-            except json.JSONDecodeError:
-                print("Failed to decode client public key from string. Aborting response.")
-                return
-        elif isinstance(client_public_key_str, dict):
-            client_public_key_jwk = client_public_key_str
-        else:
-            print("Invalid type for client public key. Expected str or dict. Aborting response.")
-            return
-        public_key = self._load_public_key_from_jwk(client_public_key_jwk)
-        if not public_key:
-            print("Failed to load public key. Aborting response.")
-            return
-        encrypted_payload = self._encrypt_with_ecdh_aesgcm(public_key, data)
-        url = f"https://{self.network}/api/tx_response/{transmitter_id}"
-        tx_response = {
-            "data": encrypted_payload,
-            "cell": self.to_dict()
-        }
-        async with aiohttp.ClientSession() as session:
-            try:
-                for _ in range(2):
-                    async with session.post(url, json=tx_response) as response:
-                        response.raise_for_status()
-                        data = await response.json()
-                print(data["message"])
-            except aiohttp.ClientError as e:
-                print(f"Error sending request: {e}")
-            except Exception as e:
-                print(f"Unexpected error: {e}")
+    def _load_env(self) -> Dict[str, str]:
+        """Load environment variables from .env file"""
+        env_path = self.config.credentials_path / ".env"
+        env_data = {}
+        
+        try:
+            with open(env_path, "r") as f:
+                for line in f:
+                    line = line.strip()
+                    if line and not line.startswith('#') and '=' in line:
+                        key, value = line.split('=', 1)
+                        env_data[key.strip()] = value.strip()
+            logger.info("Environment loaded successfully")
+            return env_data
+        except FileNotFoundError:
+            logger.error(f"Environment file not found at {env_path}")
+            return {}
+        except Exception as e:
+            logger.error(f"Error loading environment: {e}")
+            return {}
+        
+    async def __aenter__(self):
+        await self._network_client.__aenter__()
+        return self
+    
+    async def __aexit__(self, *args):
+        await self._network_client.__aexit__(*args)