Refactor API routing, add cache and security

Centralize API routing and security: moved loopback/API-key/auth/CSRF decorators into python/helpers/api.py and added register_api_route to dynamically dispatch handlers from built-in python/api and plugin api folders. Added a simple thread-safe in-memory cache (python/helpers/cache.py) to store wrapped handlers. register_api_route resolves handler classes, enforces allowed methods, composes required security wrappers, caches the resulting callables, and registers a single /api/<path> rule. Also updated runtime RFC URL to use /api/rfc. Removed the duplicate handler registration and security helper code from run_ui.py and replaced it with the new register_api_route import.

frdel committed Feb 20, 2026 at 09:43 UTC 4e243a996cb2226d9ed269fc523a9e732ee201e0
4 files changed +208 -154
python/helpers/api.py
+172 -3
@@ -1,17 +1,24 @@
1 from abc import abstractmethod
2 import json
3 +import socket
4 +import struct
5 import threading
6 +from functools import wraps
7 +from pathlib import Path
8 from typing import Union, TypedDict, Dict, Any
5 -from attr import dataclass
6 -from flask import Request, Response, jsonify, Flask, session, request, send_file
9 +from flask import Request, Response, jsonify, Flask, session, request, send_file, redirect, url_for
10 +from werkzeug.wrappers.response import Response as BaseResponse
11 from agent import AgentContext
12 from initialize import initialize_agent
13 from python.helpers.print_style import PrintStyle
14 from python.helpers.errors import format_error
11 -from werkzeug.serving import make_server
15 +from python.helpers import files, cache
16
17 ThreadLockType = Union[threading.Lock, threading.RLock]
18
19 +CACHE_AREA = "api_handlers"
20 +
21 +
22 Input = dict
23 Output = Union[Dict[str, Any], Response, TypedDict] # type: ignore
24
@@ -100,3 +107,165 @@ class ApiHandler:
107 else:
108 raise Exception(f"Context {ctxid} not found")
109
110 +
111 +
112 +
113 +def is_loopback_address(address: str) -> bool:
114 + loopback_checker = {
115 + socket.AF_INET: lambda x: (
116 + struct.unpack("!I", socket.inet_aton(x))[0] >> (32 - 8)
117 + ) == 127,
118 + socket.AF_INET6: lambda x: x == "::1",
119 + }
120 + address_type = "hostname"
121 + try:
122 + socket.inet_pton(socket.AF_INET6, address)
123 + address_type = "ipv6"
124 + except socket.error:
125 + try:
126 + socket.inet_pton(socket.AF_INET, address)
127 + address_type = "ipv4"
128 + except socket.error:
129 + address_type = "hostname"
130 +
131 + if address_type == "ipv4":
132 + return loopback_checker[socket.AF_INET](address)
133 + elif address_type == "ipv6":
134 + return loopback_checker[socket.AF_INET6](address)
135 + else:
136 + for family in (socket.AF_INET, socket.AF_INET6):
137 + try:
138 + r = socket.getaddrinfo(address, None, family, socket.SOCK_STREAM)
139 + except socket.gaierror:
140 + return False
141 + for family, _, _, _, sockaddr in r:
142 + if not loopback_checker[family](sockaddr[0]):
143 + return False
144 + return True
145 +
146 +
147 +def requires_api_key(f):
148 + @wraps(f)
149 + async def decorated(*args, **kwargs):
150 + from python.helpers.settings import get_settings
151 + valid_api_key = get_settings()["mcp_server_token"]
152 +
153 + if api_key := request.headers.get("X-API-KEY"):
154 + if api_key != valid_api_key:
155 + return Response("Invalid API key", 401)
156 + elif request.json and request.json.get("api_key"):
157 + api_key = request.json.get("api_key")
158 + if api_key != valid_api_key:
159 + return Response("Invalid API key", 401)
160 + else:
161 + return Response("API key required", 401)
162 + return await f(*args, **kwargs)
163 +
164 + return decorated
165 +
166 +
167 +def requires_loopback(f):
168 + @wraps(f)
169 + async def decorated(*args, **kwargs):
170 + if not is_loopback_address(request.remote_addr):
171 + return Response("Access denied.", 403, {})
172 + return await f(*args, **kwargs)
173 +
174 + return decorated
175 +
176 +
177 +def requires_auth(f):
178 + @wraps(f)
179 + async def decorated(*args, **kwargs):
180 + from python.helpers import login
181 + user_pass_hash = login.get_credentials_hash()
182 + if not user_pass_hash:
183 + return await f(*args, **kwargs)
184 + if session.get("authentication") != user_pass_hash:
185 + return redirect(url_for("login_handler"))
186 + return await f(*args, **kwargs)
187 +
188 + return decorated
189 +
190 +
191 +def csrf_protect(f):
192 + @wraps(f)
193 + async def decorated(*args, **kwargs):
194 + from python.helpers import runtime
195 + token = session.get("csrf_token")
196 + header = request.headers.get("X-CSRF-Token")
197 + cookie = request.cookies.get("csrf_token_" + runtime.get_runtime_id())
198 + sent = header or cookie
199 + if not token or not sent or token != sent:
200 + return Response("CSRF token missing or invalid", 403)
201 + return await f(*args, **kwargs)
202 +
203 + return decorated
204 +
205 +
206 +def register_api_route(app: Flask, lock: ThreadLockType) -> None:
207 + from python.helpers.extract_tools import load_classes_from_file
208 + from python.helpers import plugins
209 +
210 + async def _dispatch(path: str) -> BaseResponse:
211 + # Return cached wrapped handler if available
212 + cached = cache.get(CACHE_AREA, path)
213 + if cached is not None:
214 + return await cached()
215 +
216 + # Resolve file path for the handler
217 + # Try built-in api folder first, then plugin api folders
218 + handler_cls: type[ApiHandler] | None = None
219 +
220 + # Check built-in python/api/<path>.py
221 + builtin_file = files.get_abs_path(f"python/api/{path}.py")
222 + if files.is_in_dir(builtin_file, files.get_abs_path("python/api")) and files.exists(builtin_file):
223 + classes = load_classes_from_file(builtin_file, ApiHandler)
224 + if classes:
225 + handler_cls = classes[0]
226 +
227 + # Check plugin api folders: path format plugins/<plugin_name>/<handler>
228 + if handler_cls is None and path.startswith("plugins/"):
229 + parts = path.split("/", 2)
230 + if len(parts) == 3:
231 + _, plugin_name, handler_name = parts
232 + plugin_dir = plugins.find_plugin_dir(plugin_name)
233 + if plugin_dir:
234 + plugin_file = Path(plugin_dir) / "api" / f"{handler_name}.py"
235 + if plugin_file.is_file():
236 + classes = load_classes_from_file(str(plugin_file), ApiHandler)
237 + if classes:
238 + handler_cls = classes[0]
239 +
240 + if handler_cls is None:
241 + return Response(f"API endpoint not found: {path}", 404)
242 +
243 + # Check method is allowed
244 + if request.method not in handler_cls.get_methods():
245 + return Response(f"Method {request.method} not allowed for: {path}", 405)
246 +
247 + # Build handler call, wrapping with security decorators as required
248 + async def call_handler() -> BaseResponse:
249 + instance = handler_cls(app, lock)
250 + return await instance.handle_request(request=request)
251 +
252 + handler_fn = call_handler
253 + if handler_cls.requires_csrf():
254 + handler_fn = csrf_protect(handler_fn)
255 + if handler_cls.requires_api_key():
256 + handler_fn = requires_api_key(handler_fn)
257 + if handler_cls.requires_auth():
258 + handler_fn = requires_auth(handler_fn)
259 + if handler_cls.requires_loopback():
260 + handler_fn = requires_loopback(handler_fn)
261 +
262 + cache.add(CACHE_AREA, path, handler_fn)
263 + return await handler_fn()
264 +
265 + app.add_url_rule(
266 + "/api/<path:path>",
267 + "api_dispatch",
268 + _dispatch,
269 + methods=["GET", "POST", "PUT", "PATCH", "DELETE"],
270 + )
271 +
python/helpers/cache.py new
+33
@@ -0,0 +1,33 @@
1 +import threading
2 +from typing import Any
3 +
4 +_lock = threading.RLock()
5 +_cache: dict[str, dict[str, Any]] = {}
6 +
7 +
8 +def add(area: str, key: str, data: Any) -> None:
9 + with _lock:
10 + if area not in _cache:
11 + _cache[area] = {}
12 + _cache[area][key] = data
13 +
14 +
15 +def get(area: str, key: str, default: Any = None) -> Any:
16 + with _lock:
17 + return _cache.get(area, {}).get(key, default)
18 +
19 +
20 +def remove(area: str, key: str) -> None:
21 + with _lock:
22 + if area in _cache:
23 + _cache[area].pop(key, None)
24 +
25 +
26 +def clear(area: str) -> None:
27 + with _lock:
28 + _cache.pop(area, None)
29 +
30 +
31 +def clear_all() -> None:
32 + with _lock:
33 + _cache.clear()
python/helpers/runtime.py
+1 -1
@@ -138,7 +138,7 @@ def _get_rfc_url() -> str:
138 if url.endswith("/"):
139 url = url[:-1]
140 url = url + ":" + str(set["rfc_port_http"])
141 - url += "/rfc"
141 + url += "/api/rfc"
142 return url
143
144
run_ui.py
+2 -150
@@ -2,19 +2,12 @@ from datetime import timedelta
2 import os
3 import secrets
4 import time
5 -import socket
6 -import struct
7 -from functools import wraps
5 import threading
6 import asyncio
7
11 -from pathlib import Path
8 import urllib.request
13 -import urllib.error
14 -from regex.regex import F
9 import uvicorn
10 from flask import Flask, request, Response, session, redirect, url_for, render_template_string
17 -from werkzeug.wrappers.response import Response as BaseResponse
11 from werkzeug.wrappers.request import Request as WerkzeugRequest
12
13 import initialize
@@ -22,8 +15,7 @@ from python.helpers import files, git, mcp_server, fasta2a_server, settings as s
15 from python.helpers.files import get_abs_path
16 from python.helpers import runtime, dotenv, process
17 from python.helpers.websocket import WebSocketHandler, validate_ws_origin
25 -from python.helpers.extract_tools import load_classes_from_folder
26 -from python.helpers.api import ApiHandler
18 +from python.helpers.api import register_api_route, requires_auth
19 from python.helpers.print_style import PrintStyle
20 from python.helpers import login
21 import socketio # type: ignore[import-untyped]
@@ -90,107 +82,6 @@ websocket_manager.set_server_restart_broadcast(
82 # basic_auth = BasicAuth(webapp)
83
84
93 -def is_loopback_address(address):
94 - loopback_checker = {
95 - socket.AF_INET: lambda x: (
96 - struct.unpack("!I", socket.inet_aton(x))[0] >> (32 - 8)
97 - ) == 127,
98 - socket.AF_INET6: lambda x: x == "::1",
99 - }
100 - address_type = "hostname"
101 - try:
102 - socket.inet_pton(socket.AF_INET6, address)
103 - address_type = "ipv6"
104 - except socket.error:
105 - try:
106 - socket.inet_pton(socket.AF_INET, address)
107 - address_type = "ipv4"
108 - except socket.error:
109 - address_type = "hostname"
110 -
111 - if address_type == "ipv4":
112 - return loopback_checker[socket.AF_INET](address)
113 - elif address_type == "ipv6":
114 - return loopback_checker[socket.AF_INET6](address)
115 - else:
116 - for family in (socket.AF_INET, socket.AF_INET6):
117 - try:
118 - r = socket.getaddrinfo(address, None, family, socket.SOCK_STREAM)
119 - except socket.gaierror:
120 - return False
121 - for family, _, _, _, sockaddr in r:
122 - if not loopback_checker[family](sockaddr[0]):
123 - return False
124 - return True
125 -
126 -
127 -def requires_api_key(f):
128 - @wraps(f)
129 - async def decorated(*args, **kwargs):
130 - # Use the auth token from settings (same as MCP server)
131 - from python.helpers.settings import get_settings
132 - valid_api_key = get_settings()["mcp_server_token"]
133 -
134 - if api_key := request.headers.get("X-API-KEY"):
135 - if api_key != valid_api_key:
136 - return Response("Invalid API key", 401)
137 - elif request.json and request.json.get("api_key"):
138 - api_key = request.json.get("api_key")
139 - if api_key != valid_api_key:
140 - return Response("Invalid API key", 401)
141 - else:
142 - return Response("API key required", 401)
143 - return await f(*args, **kwargs)
144 -
145 - return decorated
146 -
147 -
148 -# allow only loopback addresses
149 -def requires_loopback(f):
150 - @wraps(f)
151 - async def decorated(*args, **kwargs):
152 - if not is_loopback_address(request.remote_addr):
153 - return Response(
154 - "Access denied.",
155 - 403,
156 - {},
157 - )
158 - return await f(*args, **kwargs)
159 -
160 - return decorated
161 -
162 -
163 -# require authentication for handlers
164 -def requires_auth(f):
165 - @wraps(f)
166 - async def decorated(*args, **kwargs):
167 - user_pass_hash = login.get_credentials_hash()
168 - # If no auth is configured, just proceed
169 - if not user_pass_hash:
170 - return await f(*args, **kwargs)
171 -
172 - if session.get('authentication') != user_pass_hash:
173 - return redirect(url_for('login_handler'))
174 -
175 - return await f(*args, **kwargs)
176 -
177 - return decorated
178 -
179 -
180 -def csrf_protect(f):
181 - @wraps(f)
182 - async def decorated(*args, **kwargs):
183 - token = session.get("csrf_token")
184 - header = request.headers.get("X-CSRF-Token")
185 - cookie = request.cookies.get("csrf_token_" + runtime.get_runtime_id())
186 - sent = header or cookie
187 - if not token or not sent or token != sent:
188 - return Response("CSRF token missing or invalid", 403)
189 - return await f(*args, **kwargs)
190 -
191 - return decorated
192 -
193 -
85 @webapp.route("/login", methods=["GET", "POST"])
86 async def login_handler():
87 error = None
@@ -480,46 +371,7 @@ def run():
371 runtime.get_arg("host") or dotenv.get_dotenv_value("WEB_UI_HOST") or "localhost"
372 )
373
483 - def register_api_handler(app, handler: type[ApiHandler], url_prefix: str = ""):
484 - name = handler.__module__.split(".")[-1]
485 - instance = handler(app, lock)
486 -
487 - async def handler_wrap() -> BaseResponse:
488 - return await instance.handle_request(request=request)
489 -
490 - if handler.requires_loopback():
491 - handler_wrap = requires_loopback(handler_wrap)
492 - if handler.requires_auth():
493 - handler_wrap = requires_auth(handler_wrap)
494 - if handler.requires_api_key():
495 - handler_wrap = requires_api_key(handler_wrap)
496 - if handler.requires_csrf():
497 - handler_wrap = csrf_protect(handler_wrap)
498 -
499 - route = f"{url_prefix}/{name}"
500 - app.add_url_rule(
501 - route,
502 - route,
503 - handler_wrap,
504 - methods=handler.get_methods(),
505 - )
506 -
507 - handlers = load_classes_from_folder("python/api", "*.py", ApiHandler)
508 - for handler in handlers:
509 - register_api_handler(webapp, handler, url_prefix="/api")
510 -
511 - # Load API handlers from plugins (prefixed with /plugins/{plugin_id}/)
512 - from python.helpers import plugins
513 -
514 - for plugin in plugins.get_enhanced_plugins_list():
515 - api_path = Path(plugin.path) / "api"
516 - if not api_path.exists() or not api_path.is_dir():
517 - continue
518 -
519 - plugin_handlers = load_classes_from_folder(str(api_path), "*.py", ApiHandler)
520 - for handler in plugin_handlers:
521 - # prefixed route for explicit namespacing
522 - register_api_handler(webapp, handler, url_prefix=f"/api/plugins/{plugin.name}")
374 + register_api_route(webapp, lock)
375
376 handlers_by_namespace = _build_websocket_handlers_by_namespace(socketio_server, lock)
377 configure_websocket_namespaces(