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(