Custom handlers support

pytest.ini

@@ -0,0 +1,2 @@
+[pytest]
+testpaths = tests

skyhook/handlers.py

@@ -0,0 +1,61 @@
+import inspect
+import types
+from typing import Callable, Iterator, Optional
+
+
+def skyhook_handler(fn: Optional[Callable] = None, *, html: bool = False) -> Callable:
+    """
+    Decorator to mark a function as a Skyhook handler.
+
+    Works in three forms:
+        @skyhook_handler                # no parentheses
+        @skyhook_handler()              # parentheses, no args
+        @skyhook_handler(html=True)     # with html kwarg
+
+    Sets __skyhook_handler__ = True on the function, marking it as dispatchable.
+    When html=True, also sets __skyhook_html_handler__ = True, signaling that
+    the handler returns HTML and should have origin-allowlist checks applied.
+
+    Args:
+        fn: The function being decorated (None if called with parentheses)
+        html: If True, mark this as an HTML handler
+
+    Returns:
+        The decorated function or a decorator function
+    """
+    if fn is not None and not callable(fn):
+        raise TypeError(
+            "@skyhook_handler: unexpected positional argument. "
+            "Use @skyhook_handler() or @skyhook_handler(html=True)."
+        )
+
+    def decorator(func: Callable) -> Callable:
+        func.__skyhook_handler__ = True
+        func.__skyhook_html_handler__ = html
+        return func
+
+    # Case 1: @skyhook_handler (no parentheses, fn is the decorated function)
+    if fn is not None:
+        return decorator(fn)
+
+    # Case 2 & 3: @skyhook_handler() or @skyhook_handler(html=True)
+    return decorator
+
+
+def iter_handlers(module: types.ModuleType) -> Iterator[Callable]:
+    """
+    Yield all callables in a module that have __skyhook_handler__ = True.
+
+    Skips handlers that were imported into the module rather than defined there.
+
+    Args:
+        module: A Python module to scan for handlers
+
+    Yields:
+        Callable objects marked as Skyhook handlers defined in this module
+    """
+    for name in dir(module):
+        obj = getattr(module, name)
+        if callable(obj) and getattr(obj, "__skyhook_handler__", False):
+            if inspect.getmodule(obj) is module:
+                yield obj

skyhook/responses.py

@@ -0,0 +1,28 @@
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class HtmlResponse:
+    """
+    Sentinel class for handler functions that return HTML content.
+
+    When a handler returns an HtmlResponse, the server will emit the body
+    verbatim as HTML instead of encoding it as JSON.
+    """
+    body: str
+    status: int = 200
+    content_type: str = "text/html; charset=utf-8"
+
+
+@dataclass(frozen=True)
+class RawGetResponse:
+    """
+    A raw GET response produced by a Server.raw_get_handler.
+
+    Lets an application serve arbitrary bytes (e.g. static files) for GET
+    requests that are not Skyhook function-call dispatches. The body is bytes
+    (unlike HtmlResponse.body, which is str) so any file type can be served.
+    """
+    body: bytes
+    status: int = 200
+    content_type: str = "application/octet-stream"

skyhook/server.py

@@ -8,6 +8,7 @@
 import json
 import traceback
 import urllib.parse
+from pathlib import Path
 
 from datetime import datetime
 from importlib import reload
@@ -17,6 +18,7 @@
 from .constants import Constants, Results, ServerCommands, Errors, Ports, HostPrograms, ServerEvents
 from .modules import core
 from .logger import Logger
+from .responses import HtmlResponse, RawGetResponse
 
 logger = Logger()
 
@@ -195,11 +197,15 @@ def __init__(
             load_modules: List[str] = [],
             use_main_thread_executor: bool = False,
             echo_response: bool = True,
-            only_allow_localhost_connections: bool = True
+            only_allow_localhost_connections: bool = True,
+            html_origin_allowlist_provider: Optional[Callable[[], List[str]]] = None,
+            raw_get_handler: Optional[Callable[[str], Optional[RawGetResponse]]] = None,
     ) -> None:
 
         self.events: EventEmitter = EventEmitter()
         self.executor_reply: Optional[Dict[str, Any]] = None
+        self.html_origin_allowlist_provider: Optional[Callable[[], List[str]]] = html_origin_allowlist_provider
+        self.raw_get_handler: Optional[Callable[[str], Optional[RawGetResponse]]] = raw_get_handler
 
         if port:
             self.port: int = port
@@ -347,7 +353,7 @@ def function_not_found(*args: Any, **kwargs: Any) -> None:
 
         return function_not_found
 
-    def filter_and_execute_function(self, function_name: str, parameters_dict: Dict[str, Any]) -> bytes:
+    def filter_and_execute_function(self, function_name: str, parameters_dict: Dict[str, Any]) -> Union[bytes, HtmlResponse]:
         """
         This function decides whether or the function call should come from one of the loaded modules or from the server.
         Every server function should start with SKY_
@@ -356,12 +362,16 @@ def filter_and_execute_function(self, function_name: str, parameters_dict: Dict[
 
         :param function_name: Name of the function to execute
         :param parameters_dict: Dictionary of parameters to pass to the function
-        :return: JSON response as bytes
+        :return: JSON response as bytes, or an HtmlResponse if the handler returned one
         """
         if function_name in dir(ServerCommands):
             result_json: Dict[str, Any] = self.__process_server_command(function_name, parameters_dict)
         else:
-            result_json = self.__process_module_command(function_name, parameters_dict)
+            result = self.__process_module_command(function_name, parameters_dict)
+            if isinstance(result, HtmlResponse):
+                self.executor_reply = None
+                return result
+            result_json = result
 
         self.executor_reply = None
         return json.dumps(result_json).encode()
@@ -422,7 +432,7 @@ def __process_server_command(self, function_name: str, parameters_dict: Dict[str
 
         return result_json
 
-    def __process_module_command(self, function_name: str, parameters_dict: Dict[str, Any], timeout: float = 10.0) -> Dict[str, Any]:
+    def __process_module_command(self, function_name: str, parameters_dict: Dict[str, Any], timeout: float = 10.0) -> Union[Dict[str, Any], HtmlResponse]:
         """
         Processes a command if the function in it was a module function.
 
@@ -431,7 +441,7 @@ def __process_module_command(self, function_name: str, parameters_dict: Dict[str
         :param function_name: Name of the function to execute
         :param parameters_dict: Dictionary of parameters for the function
         :param timeout: Timeout in seconds for waiting for executor reply
-        :return: Result JSON dictionary
+        :return: Result JSON dictionary, or HtmlResponse if the handler returned one
         """
         if self.__use_main_thread_executor:
             logger.debug("Emitting exec_command event")
@@ -456,6 +466,8 @@ def __process_module_command(self, function_name: str, parameters_dict: Dict[str
             return_value: Any = function(**parameters_dict)
             success: bool = True
             self.events.emit(ServerEvents.command, function_name, parameters_dict)
+            if isinstance(return_value, HtmlResponse):
+                return return_value
         except Exception as err:
             trace: str = str(traceback.format_exc())
             return_value = trace
@@ -512,19 +524,81 @@ def do_GET(self) -> None:
         if self.path == "/favicon.ico":
             return
 
+        # Static file serving — short-circuit before any command processing
+        if self.path.startswith("/static/") and self.skyhook_server.static_dir is not None:
+            filename: str = self.path[len("/static/"):]
+            file_path: Path = (self.skyhook_server.static_dir / filename).resolve()
+            static_root: Path = self.skyhook_server.static_dir.resolve()
+            if not str(file_path).startswith(str(static_root)):
+                self.send_response(403)
+                self.end_headers()
+                return
+            if not file_path.is_file():
+                self.send_response(404)
+                self.end_headers()
+                return
+            content_type, _ = mimetypes.guess_type(str(file_path))
+            if content_type is None:
+                content_type = "application/octet-stream"
+            self.send_response(200)
+            self.send_header("Content-type", content_type)
+            self.end_headers()
+            self.wfile.write(file_path.read_bytes())
+            return
+
         data: str = urllib.parse.unquote(self.path).lstrip("/")
         parts: List[str] = data.split("&")
 
         try:
             function: str = eval(parts[0])
             parameters: Dict[str, Any] = json.loads(parts[1])
-        except NameError as err:
+        except (NameError, SyntaxError, ValueError, IndexError) as err:
+            # This GET is not a function-call dispatch. Dispatch stays the primary
+            # path; only on a parse miss do we offer the request to the app's
+            # optional raw GET handler (e.g. static-file serving).
+            if self.skyhook_server.raw_get_handler is not None:
+                raw: Optional[RawGetResponse] = self.skyhook_server.raw_get_handler(self.path)
+                if raw is not None:
+                    self.send_response(raw.status)
+                    self.send_header("Content-type", raw.content_type)
+                    self.end_headers()
+                    if raw.body:
+                        self.wfile.write(raw.body)
+                    return
             logger.warning(f"Got a GET request that I don't know what to do with")
             logger.warning(f"Request was: GET {data}")
             logger.warning(f"Error is   : {err}")
             return
 
-        command_response: bytes = self.skyhook_server.filter_and_execute_function(function, parameters)
+        # Origin allowlist pre-check for HTML handlers (skip for server commands)
+        if self.skyhook_server.html_origin_allowlist_provider is not None and function not in dir(ServerCommands):
+            fn: Callable[..., Any] = self.skyhook_server.get_function_by_name(function)
+            if getattr(fn, '__skyhook_html_handler__', False):
+                origin: str = self.headers.get('Origin') or self.headers.get('Referer', '')
+                allowlist: List[str] = self.skyhook_server.html_origin_allowlist_provider()
+                if not origin:
+                    pass  # no Origin header — allow (same-origin or non-browser caller)
+                elif allowlist and not any(origin.startswith(a) for a in allowlist):
+                    self.send_response(403)
+                    self.send_header('Content-type', 'text/html; charset=utf-8')
+                    self.send_header('Access-Control-Allow-Origin', '*')
+                    self.end_headers()
+                    self.wfile.write(b"<html><body><h1>403 Forbidden</h1><p>Origin not in allowlist.</p></body></html>")
+                    return
+
+        command_response: Union[bytes, HtmlResponse] = self.skyhook_server.filter_and_execute_function(function, parameters)
+
+        # HtmlResponse dispatch — write verbatim HTML, no auto-close script
+        if isinstance(command_response, HtmlResponse):
+            self.send_response(command_response.status)
+            self.send_header('Access-Control-Allow-Origin', '*')
+            self.send_header('Access-Control-Allow-Methods', 'GET, POST, OPTIONS')
+            self.send_header('Access-Control-Allow-Headers', 'Content-Type')
+            self.send_header('Content-type', command_response.content_type)
+            self.end_headers()
+            self.wfile.write(command_response.body.encode('utf-8'))
+            return
+
         self.send_response_data("GET")
         self.wfile.write(bytes(f"{command_response}".encode("utf-8")))
         if self.reply_with_auto_close:
@@ -543,7 +617,7 @@ def do_POST(self) -> None:
         function: str = json_data.get(Constants.function_name, "")
         parameters: Dict[str, Any] = json_data.get(Constants.parameters, {})
 
-        command_response: bytes = self.skyhook_server.filter_and_execute_function(function, parameters)
+        command_response: Union[bytes, HtmlResponse] = self.skyhook_server.filter_and_execute_function(function, parameters)
         self.send_response_data("POST")
         self.wfile.write(command_response)
 

tests/test_raw_get_handler.py

@@ -0,0 +1,101 @@
+"""Tests for the generic raw_get_handler GET-fallback seam in the server."""
+import socket
+import threading
+import time
+import urllib.error
+import urllib.parse
+import urllib.request
+
+import pytest
+
+from skyhook.responses import RawGetResponse
+from skyhook.server import Server
+
+
+def _find_free_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("127.0.0.1", 0))
+        return s.getsockname()[1]
+
+
+def _wait_for_port(port: int, host: str = "127.0.0.1", timeout: float = 5.0) -> None:
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        try:
+            with socket.create_connection((host, port), timeout=0.1):
+                return
+        except OSError:
+            time.sleep(0.05)
+    raise TimeoutError(f"Server on port {port} did not start within {timeout}s")
+
+
+class ServerContext:
+    def __init__(self, port: int, **server_kwargs):
+        self.port = port
+        self._kwargs = server_kwargs
+        self.server = None
+        self._thread = None
+
+    def __enter__(self):
+        self.server = Server(port=self.port, echo_response=False, **self._kwargs)
+        self._thread = threading.Thread(target=self.server.start_listening, daemon=True)
+        self._thread.start()
+        _wait_for_port(self.port)
+        return self
+
+    def __exit__(self, *_):
+        if self.server is not None:
+            self.server.stop_listening()
+        if self._thread is not None:
+            self._thread.join(timeout=3)
+
+
+def test_raw_get_handler_serves_non_dispatch_get():
+    calls = []
+
+    def handler(path):
+        calls.append(path)
+        return RawGetResponse(b"hello raw", status=200, content_type="text/plain")
+
+    port = _find_free_port()
+    with ServerContext(port, raw_get_handler=handler) as ctx:
+        resp = urllib.request.urlopen(f"http://127.0.0.1:{ctx.port}/some/raw/path", timeout=5)
+        assert resp.status == 200
+        assert resp.headers.get("Content-type") == "text/plain"
+        assert resp.read() == b"hello raw"
+    assert calls == ["/some/raw/path"]
+
+
+def test_dispatch_call_not_routed_to_raw_get_handler():
+    calls = []
+
+    def handler(path):
+        calls.append(path)
+        return RawGetResponse(b"should not happen", status=200)
+
+    port = _find_free_port()
+    with ServerContext(port, raw_get_handler=handler) as ctx:
+        # A well-formed dispatch call for an unknown function still PARSES, so it
+        # goes through dispatch (returns function_not_found JSON), never the seam.
+        encoded = "/" + urllib.parse.quote('"nonexistent_fn"&{}', safe="")
+        try:
+            urllib.request.urlopen(f"http://127.0.0.1:{ctx.port}{encoded}", timeout=5)
+        except Exception:
+            pass
+    assert calls == []
+
+
+def test_raw_get_handler_none_falls_through():
+    calls = []
+
+    def handler(path):
+        calls.append(path)
+        return None
+
+    port = _find_free_port()
+    with ServerContext(port, raw_get_handler=handler) as ctx:
+        try:
+            urllib.request.urlopen(f"http://127.0.0.1:{ctx.port}/definitely/not/a/call", timeout=5)
+        except Exception:
+            pass
+    assert calls == ["/definitely/not/a/call"]