hypern 0.3.0__cp311-cp311-win32.whl → 0.3.2__cp311-cp311-win32.whl

hypern/middleware/security.py

@@ -0,0 +1,184 @@
+import hashlib
+import hmac
+import secrets
+import time
+from base64 import b64decode, b64encode
+from dataclasses import dataclass
+from datetime import datetime, timedelta, timezone
+from typing import Any, Dict, List, Optional
+
+import jwt
+
+from hypern.exceptions import Forbidden, Unauthorized
+from hypern.hypern import Request, Response
+from .base import Middleware, MiddlewareConfig
+
+
+@dataclass
+class CORSConfig:
+    allowed_origins: List[str]
+    allowed_methods: List[str]
+    max_age: int
+
+
+@dataclass
+class SecurityConfig:
+    rate_limiting: bool = False
+    jwt_auth: bool = False
+    cors_configuration: Optional[CORSConfig] = None
+    csrf_protection: bool = False
+    security_headers: Optional[Dict[str, str]] = None
+    jwt_secret: str = ""
+    jwt_algorithm: str = "HS256"
+    jwt_expires_in: int = 3600  # 1 hour in seconds
+
+    def __post_init__(self):
+        if self.cors_configuration:
+            self.cors_configuration = CORSConfig(**self.cors_configuration)
+
+        if self.security_headers is None:
+            self.security_headers = {
+                "X-Frame-Options": "DENY",
+                "X-Content-Type-Options": "nosniff",
+                "Strict-Transport-Security": "max-age=31536000; includeSubDomains",
+            }
+
+
+class SecurityMiddleware(Middleware):
+    def __init__(self, secur_config: SecurityConfig, config: Optional[MiddlewareConfig] = None):
+        super().__init__(config)
+        self.secur_config = secur_config
+        self._secret_key = secrets.token_bytes(32)
+        self._token_lifetime = 3600
+        self._rate_limit_storage = {}
+
+    def _rate_limit_check(self, request: Request) -> Optional[Response]:
+        """Check if the request exceeds rate limits"""
+        if not self.secur_config.rate_limiting:
+            return None
+
+        client_ip = request.client.host
+        current_time = time.time()
+        window_start = int(current_time - 60)  # 1-minute window
+
+        # Clean up old entries
+        self._rate_limit_storage = {ip: hits for ip, hits in self._rate_limit_storage.items() if hits["timestamp"] > window_start}
+
+        if client_ip not in self._rate_limit_storage:
+            self._rate_limit_storage[client_ip] = {"count": 1, "timestamp": current_time}
+        else:
+            self._rate_limit_storage[client_ip]["count"] += 1
+
+        if self._rate_limit_storage[client_ip]["count"] > 60:  # 60 requests per minute
+            return Response(status_code=429, description=b"Too Many Requests", headers={"Retry-After": "60"})
+        return None
+
+    def _generate_jwt_token(self, user_data: Dict[str, Any]) -> str:
+        """Generate a JWT token"""
+        if not self.secur_config.jwt_secret:
+            raise ValueError("JWT secret key is not configured")
+
+        payload = {
+            "user": user_data,
+            "exp": datetime.now(tz=timezone.utc) + timedelta(seconds=self.secur_config.jwt_expires_in),
+            "iat": datetime.now(tz=timezone.utc),
+        }
+        return jwt.encode(payload, self.secur_config.jwt_secret, algorithm=self.secur_config.jwt_algorithm)
+
+    def _verify_jwt_token(self, token: str) -> Dict[str, Any]:
+        """Verify JWT token and return payload"""
+        try:
+            payload = jwt.decode(token, self.secur_config.jwt_secret, algorithms=[self.secur_config.jwt_algorithm])
+            return payload
+        except jwt.ExpiredSignatureError:
+            raise Unauthorized("Token has expired")
+        except jwt.InvalidTokenError:
+            raise Unauthorized("Invalid token")
+
+    def _generate_csrf_token(self, session_id: str) -> str:
+        """Generate a new CSRF token"""
+        timestamp = str(int(time.time()))
+        token_data = f"{session_id}:{timestamp}"
+        signature = hmac.new(self._secret_key, token_data.encode(), hashlib.sha256).digest()
+        return b64encode(f"{token_data}:{b64encode(signature).decode()}".encode()).decode()
+
+    def _validate_csrf_token(self, token: str) -> bool:
+        """Validate CSRF token"""
+        try:
+            decoded_token = b64decode(token.encode()).decode()
+            session_id, timestamp, signature = decoded_token.rsplit(":", 2)
+
+            # Verify timestamp
+            token_time = int(timestamp)
+            current_time = int(time.time())
+            if current_time - token_time > self._token_lifetime:
+                return False
+
+            # Verify signature
+            expected_data = f"{session_id}:{timestamp}"
+            expected_signature = hmac.new(self._secret_key, expected_data.encode(), hashlib.sha256).digest()
+
+            actual_signature = b64decode(signature)
+            return hmac.compare_digest(expected_signature, actual_signature)
+
+        except (ValueError, AttributeError, TypeError):
+            return False
+
+    def _apply_cors_headers(self, response: Response) -> None:
+        """Apply CORS headers to response"""
+        if not self.secur_config.cors_configuration:
+            return
+
+        cors = self.secur_config.cors_configuration
+        response.headers.update(
+            {
+                "Access-Control-Allow-Origin": ", ".join(cors.allowed_origins),
+                "Access-Control-Allow-Methods": ", ".join(cors.allowed_methods),
+                "Access-Control-Max-Age": str(cors.max_age),
+                "Access-Control-Allow-Headers": "Content-Type, Authorization, X-CSRF-Token",
+                "Access-Control-Allow-Credentials": "true",
+            }
+        )
+
+    def _apply_security_headers(self, response: Response) -> None:
+        """Apply security headers to response"""
+        if self.secur_config.security_headers:
+            response.headers.update(self.secur_config.security_headers)
+
+    async def before_request(self, request: Request) -> Request | Response:
+        """Process request before handling"""
+        # Rate limiting check
+        if rate_limit_response := self._rate_limit_check(request):
+            return rate_limit_response
+
+        # JWT authentication check
+        if self.secur_config.jwt_auth:
+            auth_header = request.headers.get("Authorization")
+            if not auth_header or not auth_header.startswith("Bearer "):
+                raise Unauthorized("Missing or invalid authorization header")
+            token = auth_header.split(" ")[1]
+            try:
+                request.user = self._verify_jwt_token(token)
+            except Unauthorized as e:
+                return Response(status_code=401, description=str(e))
+
+        # CSRF protection check
+        if self.secur_config.csrf_protection and request.method in ["POST", "PUT", "DELETE", "PATCH"]:
+            csrf_token = request.headers.get("X-CSRF-Token")
+            if not csrf_token or not self._validate_csrf_token(csrf_token):
+                raise Forbidden("CSRF token missing or invalid")
+
+        return request
+
+    async def after_request(self, response: Response) -> Response:
+        """Process response after handling"""
+        self._apply_security_headers(response)
+        self._apply_cors_headers(response)
+        return response
+
+    def generate_csrf_token(self, request: Request) -> str:
+        """Generate and set CSRF token for the request"""
+        if not hasattr(request, "session_id"):
+            request.session_id = secrets.token_urlsafe(32)
+        token = self._generate_csrf_token(request.session_id)
+        return token

hypern/processpool.py

@@ -25,11 +25,26 @@ def run_processes(
     after_request: List[FunctionInfo],
     response_headers: Dict[str, str],
     reload: bool = True,
+    on_startup: FunctionInfo | None = None,
+    on_shutdown: FunctionInfo | None = None,
+    auto_compression: bool = False,
 ) -> List[Process]:
     socket = SocketHeld(host, port)
 
     process_pool = init_processpool(
-        router, websocket_router, socket, workers, processes, max_blocking_threads, injectables, before_request, after_request, response_headers
+        router,
+        websocket_router,
+        socket,
+        workers,
+        processes,
+        max_blocking_threads,
+        injectables,
+        before_request,
+        after_request,
+        response_headers,
+        on_startup,
+        on_shutdown,
+        auto_compression,
     )
 
     def terminating_signal_handler(_sig, _frame):
@@ -79,6 +94,9 @@ def init_processpool(
     before_request: List[FunctionInfo],
     after_request: List[FunctionInfo],
     response_headers: Dict[str, str],
+    on_startup: FunctionInfo | None = None,
+    on_shutdown: FunctionInfo | None = None,
+    auto_compression: bool = False,
 ) -> List[Process]:
     process_pool = []
 
@@ -86,7 +104,20 @@ def init_processpool(
         copied_socket = socket.try_clone()
         process = Process(
             target=spawn_process,
-            args=(router, websocket_router, copied_socket, workers, max_blocking_threads, injectables, before_request, after_request, response_headers),
+            args=(
+                router,
+                websocket_router,
+                copied_socket,
+                workers,
+                max_blocking_threads,
+                injectables,
+                before_request,
+                after_request,
+                response_headers,
+                on_startup,
+                on_shutdown,
+                auto_compression,
+            ),
         )
         process.start()
         process_pool.append(process)
@@ -118,6 +149,9 @@ def spawn_process(
     before_request: List[FunctionInfo],
     after_request: List[FunctionInfo],
     response_headers: Dict[str, str],
+    on_startup: FunctionInfo | None = None,
+    on_shutdown: FunctionInfo | None = None,
+    auto_compression: bool = False,
 ):
     loop = initialize_event_loop()
 
@@ -128,7 +162,12 @@ def spawn_process(
     server.set_before_hooks(hooks=before_request)
     server.set_after_hooks(hooks=after_request)
     server.set_response_headers(headers=response_headers)
+    server.set_auto_compression(enabled=auto_compression)
 
+    if on_startup:
+        server.set_startup_handler(on_startup)
+    if on_shutdown:
+        server.set_shutdown_handler(on_shutdown)
     try:
         server.start(socket, workers, max_blocking_threads)
         loop = asyncio.get_event_loop()

hypern/reload.py

@@ -2,6 +2,8 @@ import sys
 import time
 import subprocess
 from watchdog.events import FileSystemEventHandler
+import signal
+import os
 
 from .logging import logger
 
@@ -10,51 +12,35 @@ class EventHandler(FileSystemEventHandler):
     def __init__(self, file_path: str, directory_path: str) -> None:
         self.file_path = file_path
         self.directory_path = directory_path
-        self.process = None  # Keep track of the subprocess
-        self.last_reload = time.time()  # Keep track of the last reload. EventHandler is initialized with the process.
-
-    def stop_server(self):
-        if self.process:
-            try:
-                # Check if the process is still alive
-                if self.process.poll() is None:  # None means the process is still running
-                    self.process.terminate()  # Gracefully terminate the process
-                    self.process.wait(timeout=5)  # Wait for the process to exit
-                else:
-                    logger.error("Process is not running.")
-            except subprocess.TimeoutExpired:
-                logger.error("Process did not terminate in time. Forcing termination.")
-                self.process.kill()  # Forcefully kill the process if it doesn't stop
-            except ProcessLookupError:
-                logger.error("Process does not exist.")
-            except Exception as e:
-                logger.error(f"An error occurred while stopping the server: {e}")
-        else:
-            logger.debug("No process to stop.")
+        self.process = None
+        self.last_reload = time.time()
 
     def reload(self):
-        self.stop_server()
-        logger.debug("Reloading the server")
-        prev_process = self.process
-        if prev_process:
-            prev_process.kill()
-
-        self.process = subprocess.Popen(
-            [sys.executable, *sys.argv],
-        )
-
-        self.last_reload = time.time()
+        # Kill all existing processes with the same command
+        current_cmd = [sys.executable, *sys.argv]
+
+        try:
+            # Find and kill existing processes
+            for proc in subprocess.Popen(["ps", "aux"], stdout=subprocess.PIPE).communicate()[0].decode().splitlines():
+                if all(str(arg) in proc for arg in current_cmd):
+                    pid = int(proc.split()[1])
+                    try:
+                        os.kill(pid, signal.SIGKILL)  # NOSONAR
+                        logger.debug(f"Killed process with PID {pid}")
+                    except ProcessLookupError:
+                        pass
+
+            # Start new process
+            self.process = subprocess.Popen(current_cmd)
+            self.last_reload = time.time()
+            logger.debug("Server reloaded successfully")
+
+        except Exception as e:
+            logger.error(f"Reload failed: {e}")
 
     def on_modified(self, event) -> None:
-        """
-        This function is a callback that will start a new server on every even change
-
-        :param event FSEvent: a data structure with info about the events
-        """
-
-        # Avoid reloading multiple times when watchdog detects multiple events
         if time.time() - self.last_reload < 0.5:
             return
 
-        time.sleep(0.2)  # Wait for the file to be fully written
+        time.sleep(0.2)  # Ensure file is written
         self.reload()

hypern/routing/__init__.py

@@ -1,4 +1,5 @@
 from .route import Route
 from .endpoint import HTTPEndpoint
+from .queue import QueuedHTTPEndpoint
 
-__all__ = ["Route", "HTTPEndpoint"]
+__all__ = ["Route", "HTTPEndpoint", "QueuedHTTPEndpoint"]

hypern/routing/queue.py

@@ -0,0 +1,175 @@
+import asyncio
+import time
+from contextlib import asynccontextmanager
+from dataclasses import dataclass, field
+from queue import PriorityQueue
+from typing import Any, Dict
+
+from hypern.hypern import Request, Response
+from hypern.response import JSONResponse
+from hypern.routing import HTTPEndpoint
+from hypern.logging import logger
+
+
+@dataclass(order=True)
+class PrioritizedRequest:
+    priority: int
+    timestamp: float = field(default_factory=time.time)
+    request: Request | None = field(default=None, compare=False)
+    future: asyncio.Future = field(compare=False, default_factory=asyncio.Future)
+
+
+class QueuedHTTPEndpoint(HTTPEndpoint):
+    """
+    HTTPEndpoint with request queuing capabilities for high-load scenarios.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # Queue configuration
+        self._max_concurrent = kwargs.get("max_concurrent", 100)
+        self._queue_size = kwargs.get("queue_size", 1000)
+        self._request_timeout = kwargs.get("request_timeout", 30)
+
+        # Initialize queuing system
+        self._request_queue: PriorityQueue = PriorityQueue(maxsize=self._queue_size)
+        self._active_requests = 0
+        self._lock = None  # Will be initialized when needed
+        self._request_semaphore = None  # Will be initialized when needed
+        self._shutdown = False
+        self._queue_task = None
+        self._initialized = False
+
+        # Metrics
+        self._metrics = {"processed_requests": 0, "queued_requests": 0, "rejected_requests": 0, "avg_wait_time": 0.0}
+
+        self._fully_message = "Request queue is full"
+
+    async def _initialize(self):
+        """Initialize async components when first request arrives"""
+        if not self._initialized:
+            self._lock = asyncio.Lock()
+            self._request_semaphore = asyncio.Semaphore(self._max_concurrent)
+            self._queue_task = asyncio.create_task(self._process_queue())
+            self._initialized = True
+
+    @asynccontextmanager
+    async def _queue_context(self, request: Request, priority: int = 10):
+        """Context manager for handling request queuing."""
+        if self._shutdown:
+            raise RuntimeError("Endpoint is shutting down")
+
+        await self._initialize()  # Ensure async components are initialized
+
+        request_future = asyncio.Future()
+        prioritized_request = PrioritizedRequest(priority=priority, timestamp=time.time(), request=request, future=request_future)
+
+        try:
+            if self._request_queue.qsize() >= self._queue_size:
+                self._metrics["rejected_requests"] += 1
+                raise asyncio.QueueFull(self._fully_message)
+
+            await self._enqueue_request(prioritized_request)
+            yield await asyncio.wait_for(request_future, timeout=self._request_timeout)
+
+        except asyncio.TimeoutError:
+            self._metrics["rejected_requests"] += 1
+            raise asyncio.TimeoutError("Request timed out while waiting in queue")
+        finally:
+            if not request_future.done():
+                request_future.cancel()
+
+    async def _enqueue_request(self, request: PrioritizedRequest):
+        """Add request to the queue."""
+        try:
+            self._request_queue.put_nowait(request)
+            self._metrics["queued_requests"] += 1
+        except asyncio.QueueFull:
+            self._metrics["rejected_requests"] += 1
+            raise asyncio.QueueFull(self._fully_message)
+
+    async def _process_queue(self):
+        """Background task to process queued requests."""
+        while not self._shutdown:
+            try:
+                if not self._request_queue.empty():
+                    async with self._lock:
+                        if self._active_requests >= self._max_concurrent:
+                            await asyncio.sleep(0.1)
+                            continue
+
+                        request = self._request_queue.get_nowait()
+                        wait_time = time.time() - request.timestamp
+                        self._metrics["avg_wait_time"] = (self._metrics["avg_wait_time"] * self._metrics["processed_requests"] + wait_time) / (
+                            self._metrics["processed_requests"] + 1
+                        )
+
+                        if not request.future.cancelled():
+                            self._active_requests += 1
+                            asyncio.create_task(self._handle_request(request))
+
+                await asyncio.sleep(0.01)
+            except Exception as e:
+                logger.error(f"Error processing queue: {e}")
+                await asyncio.sleep(1)
+
+    async def _handle_request(self, request: PrioritizedRequest):
+        """Handle individual request."""
+        try:
+            async with self._request_semaphore:
+                response = await super().dispatch(request.request, {})
+                if not request.future.done():
+                    request.future.set_result(response)
+        except Exception as e:
+            if not request.future.done():
+                request.future.set_exception(e)
+        finally:
+            self._active_requests -= 1
+            self._metrics["processed_requests"] += 1
+            self._metrics["queued_requests"] -= 1
+
+    async def dispatch(self, request: Request, inject: Dict[str, Any]) -> Response:
+        """
+        Enhanced dispatch method with request queuing.
+        """
+        try:
+            priority = self._get_request_priority(request)
+
+            async with self._queue_context(request, priority) as response:
+                return response
+
+        except asyncio.QueueFull:
+            return JSONResponse(description={"error": "Server too busy", "message": self._fully_message, "retry_after": 5}, status_code=503)
+        except asyncio.TimeoutError:
+            return JSONResponse(
+                description={
+                    "error": "Request timeout",
+                    "message": "Request timed out while waiting in queue",
+                },
+                status_code=504,
+            )
+        except Exception as e:
+            return JSONResponse(description={"error": "Internal server error", "message": str(e)}, status_code=500)
+
+    def _get_request_priority(self, request: Request) -> int:
+        """
+        Determine request priority. Override this method to implement
+        custom priority logic.
+        """
+        if request.method == "GET":
+            return 5
+        return 10
+
+    async def shutdown(self):
+        """Gracefully shutdown the endpoint."""
+        self._shutdown = True
+        if self._queue_task and not self._queue_task.done():
+            self._queue_task.cancel()
+            try:
+                await self._queue_task
+            except asyncio.CancelledError:
+                pass
+
+    def get_metrics(self) -> Dict[str, Any]:
+        """Get current queue metrics."""
+        return {**self._metrics, "current_queue_size": self._request_queue.qsize(), "active_requests": self._active_requests}

hypern/ws/__init__.py

@@ -0,0 +1,4 @@
+from hypern.hypern import WebSocketSession
+from .route import WebsocketRoute
+
+__all__ = ["WebsocketRoute", "WebSocketSession"]

hypern/ws/channel.py

@@ -0,0 +1,80 @@
+from dataclasses import dataclass, field
+from typing import Any, Awaitable, Callable, Dict, Set
+
+from hypern.hypern import WebSocketSession
+
+
+@dataclass
+class Channel:
+    name: str
+    subscribers: Set[WebSocketSession] = field(default_factory=set)
+    handlers: Dict[str, Callable[[WebSocketSession, Any], Awaitable[None]]] = field(default_factory=dict)
+
+    def publish(self, event: str, data: Any, publisher: WebSocketSession = None):
+        """Publish an event to all subscribers except the publisher"""
+        for subscriber in self.subscribers:
+            if subscriber != publisher:
+                subscriber.send({"channel": self.name, "event": event, "data": data})
+
+    def handle_event(self, event: str, session: WebSocketSession, data: Any):
+        """Handle an event on this channel"""
+        if event in self.handlers:
+            self.handlers[event](session, data)
+
+    def add_subscriber(self, subscriber: WebSocketSession):
+        """Add a subscriber to the channel"""
+        self.subscribers.add(subscriber)
+
+    def remove_subscriber(self, subscriber: WebSocketSession):
+        """Remove a subscriber from the channel"""
+        self.subscribers.discard(subscriber)
+
+    def on(self, event: str):
+        """Decorator for registering event handlers"""
+
+        def decorator(handler: Callable[[WebSocketSession, Any], Awaitable[None]]):
+            self.handlers[event] = handler
+            return handler
+
+        return decorator
+
+
+class ChannelManager:
+    def __init__(self):
+        self.channels: Dict[str, Channel] = {}
+        self.client_channels: Dict[WebSocketSession, Set[str]] = {}
+
+    def create_channel(self, channel_name: str) -> Channel:
+        """Create a new channel if it doesn't exist"""
+        if channel_name not in self.channels:
+            self.channels[channel_name] = Channel(channel_name)
+        return self.channels[channel_name]
+
+    def get_channel(self, channel_name: str) -> Channel:
+        """Get a channel by name"""
+        return self.channels.get(channel_name)
+
+    def subscribe(self, client: WebSocketSession, channel_name: str):
+        """Subscribe a client to a channel"""
+        channel = self.create_channel(channel_name)
+        channel.add_subscriber(client)
+
+        if client not in self.client_channels:
+            self.client_channels[client] = set()
+        self.client_channels[client].add(channel_name)
+
+    def unsubscribe(self, client: WebSocketSession, channel_name: str):
+        """Unsubscribe a client from a channel"""
+        channel = self.get_channel(channel_name)
+        if channel:
+            channel.remove_subscriber(client)
+            if client in self.client_channels:
+                self.client_channels[client].discard(channel_name)
+
+    def unsubscribe_all(self, client: WebSocketSession):
+        """Unsubscribe a client from all channels"""
+        if client in self.client_channels:
+            channels = self.client_channels[client].copy()
+            for channel_name in channels:
+                self.unsubscribe(client, channel_name)
+            del self.client_channels[client]