main
py 302 lines 9.86 KB
Raw
1 from abc import abstractmethod
2 import json
3 import threading
4 from urllib.parse import urlsplit, unquote
5 from functools import wraps
6 from pathlib import Path
7 from typing import Union, Dict, Any
8 from flask import (
9 Request,
10 Response,
11 jsonify,
12 Flask,
13 session,
14 request,
15 send_file,
16 redirect,
17 url_for,
18 )
19 from werkzeug.wrappers.response import Response as BaseResponse
20 from helpers.print_style import PrintStyle
21 from helpers.errors import format_error
22 from helpers import files, cache
23
24 ThreadLockType = Union[threading.Lock, threading.RLock]
25
26 CACHE_AREA = "api_handlers(api)"
27 # cache.toggle_area(CACHE_AREA, False) # cache off for now
28
29 Input = dict
30 Output = Union[Dict[str, Any], Response]
31
32
33 class ApiHandler:
34 def __init__(self, app: Flask, thread_lock: ThreadLockType):
35 self.app = app
36 self.thread_lock = thread_lock
37
38 @classmethod
39 def requires_loopback(cls) -> bool:
40 return False
41
42 @classmethod
43 def requires_api_key(cls) -> bool:
44 return False
45
46 @classmethod
47 def requires_auth(cls) -> bool:
48 return True
49
50 @classmethod
51 def get_methods(cls) -> list[str]:
52 return ["POST"]
53
54 @classmethod
55 def requires_csrf(cls) -> bool:
56 return cls.requires_auth()
57
58 @abstractmethod
59 async def process(self, input: Input, request: Request) -> Output:
60 pass
61
62 async def handle_request(self, request: Request) -> Response:
63 try:
64 # input data from request based on type
65 input_data: Input = {}
66 if request.is_json:
67 try:
68 if request.data: # Check if there's any data
69 input_data = request.get_json()
70 # If empty or not valid JSON, use empty dict
71 except Exception as e:
72 # Just log the error and continue with empty input
73 PrintStyle().print(f"Error parsing JSON: {str(e)}")
74 input_data = {}
75 else:
76 # input_data = {"data": request.get_data(as_text=True)}
77 input_data = {}
78
79 # process via handler
80 output = await self.process(input_data, request)
81
82 # return output based on type
83 if isinstance(output, Response):
84 return output
85 else:
86 response_json = json.dumps(output)
87 return Response(
88 response=response_json, status=200, mimetype="application/json"
89 )
90
91 # return exceptions with 500
92 except Exception as e:
93 error = format_error(e)
94 PrintStyle.error(f"API error: {error}")
95 return Response(response=error, status=500, mimetype="text/plain")
96
97 # get context to run agent zero in
98 def use_context(self, ctxid: str, create_if_not_exists: bool = True):
99 from helpers.context_utils import use_context as _use_context
100 return _use_context(self.thread_lock, ctxid, create_if_not_exists)
101
102
103 from helpers.network import is_loopback_address
104
105
106 def is_safe_next_url(value: str | None) -> bool:
107 """Return True when value is a safe same-origin redirect target."""
108 if not value:
109 return False
110 if "\r" in value or "\n" in value:
111 return False
112 # Reject raw backslashes (browsers normalize `/\host` to `//host` -> external).
113 if "\\" in value:
114 return False
115
116 # Decode percent-escapes so encoded backslashes (e.g. `%5C`) are caught too.
117 decoded = unquote(value)
118 if "\\" in decoded:
119 return False
120
121 parsed = urlsplit(decoded)
122 if parsed.scheme or parsed.netloc:
123 return False
124
125 # Require an absolute path within this origin, but reject protocol-relative URLs.
126 return parsed.path.startswith("/") and not parsed.path.startswith("//")
127
128
129 def get_safe_next_url(value: str | None, fallback: str | None = None) -> str | None:
130 """Return value if it is a safe next URL, otherwise return a safe fallback."""
131 if is_safe_next_url(value):
132 return value
133 if is_safe_next_url(fallback):
134 return fallback
135 return None
136
137
138 def get_current_request_next_url() -> str:
139 """Return the current request path/query as a safe relative redirect target."""
140 next_url = request.full_path if request.query_string else request.path
141 return get_safe_next_url(next_url, url_for("serve_index")) or url_for("serve_index")
142
143
144 def requires_api_key(f):
145 @wraps(f)
146 async def decorated(*args, **kwargs):
147 from helpers.settings import get_settings
148
149 valid_api_key = get_settings()["mcp_server_token"]
150
151 if api_key := request.headers.get("X-API-KEY"):
152 if api_key != valid_api_key:
153 return Response("Invalid API key", 401)
154 elif request.json and request.json.get("api_key"):
155 api_key = request.json.get("api_key")
156 if api_key != valid_api_key:
157 return Response("Invalid API key", 401)
158 else:
159 return Response("API key required", 401)
160 return await f(*args, **kwargs)
161
162 return decorated
163
164
165 def requires_loopback(f):
166 @wraps(f)
167 async def decorated(*args, **kwargs):
168 if not is_loopback_address(str(request.remote_addr)):
169 return Response("Access denied.", 403, {})
170 return await f(*args, **kwargs)
171
172 return decorated
173
174
175 def requires_auth(f):
176 @wraps(f)
177 async def decorated(*args, **kwargs):
178 from helpers import login
179
180 user_pass_hash = login.get_credentials_hash()
181 if not user_pass_hash:
182 return await f(*args, **kwargs)
183 if session.get("authentication") != user_pass_hash:
184 return redirect(url_for("login_handler", next=get_current_request_next_url()))
185 return await f(*args, **kwargs)
186
187 return decorated
188
189
190 def csrf_protect(f):
191 @wraps(f)
192 async def decorated(*args, **kwargs):
193 from helpers import runtime
194
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 helpers.modules import load_classes_from_file
208 from 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 and plugin api folders before the user fallback
218 handler_cls: type[ApiHandler] | None = None
219
220 # Check built-in python/api/<path>.py
221 builtin_file = files.get_abs_path(f"api/{path}.py")
222 if files.is_in_dir(builtin_file, files.get_abs_path("api")) and files.exists(
223 builtin_file
224 ):
225 classes = load_classes_from_file(builtin_file, ApiHandler)
226 if classes:
227 handler_cls = classes[0]
228
229 # Check plugin api folders: path format plugins/<plugin_name>/<handler>
230 if handler_cls is None and path.startswith("plugins/"):
231 parts = path.split("/", 2)
232 if len(parts) == 3:
233 _, plugin_name, handler_name = parts
234 plugin_dir = plugins.find_plugin_dir(plugin_name)
235 if plugin_dir:
236 plugin_file = Path(plugin_dir) / "api" / f"{handler_name}.py"
237 if plugin_file.is_file():
238 classes = load_classes_from_file(str(plugin_file), ApiHandler)
239 if classes:
240 handler_cls = classes[0]
241
242 # Check user api/<path>.py
243 if handler_cls is None:
244 user_api_dir = files.get_abs_path(files.USER_DIR, files.API_DIR)
245 user_file = files.get_abs_path(user_api_dir, f"{path}.py")
246 if files.is_in_dir(user_file, user_api_dir) and files.exists(user_file):
247 classes = load_classes_from_file(user_file, ApiHandler)
248 if classes:
249 handler_cls = classes[0]
250
251 if handler_cls is None:
252 return Response(f"API endpoint not found: {path}", 404)
253
254 # Check method is allowed
255 if request.method not in handler_cls.get_methods():
256 return Response(f"Method {request.method} not allowed for: {path}", 405)
257
258 # Build handler call, wrapping with security decorators as required
259 async def call_handler() -> BaseResponse:
260 instance = handler_cls(app, lock)
261 return await instance.handle_request(request=request)
262
263 handler_fn = call_handler
264 if handler_cls.requires_csrf():
265 handler_fn = csrf_protect(handler_fn)
266 if handler_cls.requires_api_key():
267 handler_fn = requires_api_key(handler_fn)
268 if handler_cls.requires_auth():
269 handler_fn = requires_auth(handler_fn)
270 if handler_cls.requires_loopback():
271 handler_fn = requires_loopback(handler_fn)
272
273 cache.add(CACHE_AREA, path, handler_fn)
274 return await handler_fn()
275
276 app.add_url_rule(
277 "/api/<path:path>",
278 "api_dispatch",
279 _dispatch,
280 methods=["GET", "POST", "PUT", "PATCH", "DELETE"],
281 )
282
283
284 def register_watchdogs():
285 from helpers import watchdog
286 from helpers.ws import CACHE_AREA as WS_CACHE_AREA
287
288
289 def on_api_change(items: list[watchdog.WatchItem]):
290 PrintStyle.debug("API endpoint watchdog triggered:", items)
291 cache.clear(CACHE_AREA)
292 cache.clear(WS_CACHE_AREA)
293
294 watchdog.add_watchdog(
295 "api_handlers",
296 roots=[
297 files.get_abs_path(files.API_DIR),
298 files.get_abs_path(files.USER_DIR, files.API_DIR),
299 ],
300 patterns=["*.py"],
301 handler=on_api_change,
302 )