refactor: stateless auth and parser hardening

This commit is contained in:
cra88y
2026-01-11 09:46:45 -06:00
parent ea53d8dceb
commit daacae3842
2 changed files with 25 additions and 20 deletions
+16 -17
View File
@@ -1,4 +1,4 @@
import json, http.client, ssl, os, logging, base64, secrets, time import json, http.client, ssl, os, logging, base64, secrets, time, hashlib
from flask import Response, request, abort from flask import Response, request, abort
from searx.plugins import Plugin, PluginInfo from searx.plugins import Plugin, PluginInfo
from searx.result_types import EngineResults from searx.result_types import EngineResults
@@ -24,25 +24,24 @@ class SXNGPlugin(Plugin):
self.max_tokens = int(os.getenv('GEMINI_MAX_TOKENS', 500)) self.max_tokens = int(os.getenv('GEMINI_MAX_TOKENS', 500))
self.temperature = float(os.getenv('GEMINI_TEMPERATURE', 0.2)) self.temperature = float(os.getenv('GEMINI_TEMPERATURE', 0.2))
self.base_url = os.getenv('OPENROUTER_BASE_URL', 'openrouter.ai') self.base_url = os.getenv('OPENROUTER_BASE_URL', 'openrouter.ai')
self.valid_tokens = {} self.secret = os.getenv('SXNG_LLM_SECRET', secrets.token_hex(32))
def init(self, app): def init(self, app):
@app.route('/gemini-stream', methods=['POST']) @app.route('/gemini-stream', methods=['POST'])
def g_stream(): def g_stream():
data = request.json or {} data = request.json or {}
token = data.get('tk', '') token = data.get('tk', '')
# Maintenance: Token validation & cleanup
now = time.time()
self.valid_tokens = {k: v for k, v in self.valid_tokens.items() if v > now}
if token not in self.valid_tokens:
abort(403)
del self.valid_tokens[token]
context_text = data.get('context', '')
q = data.get('q', '') q = data.get('q', '')
try:
ts, sig = token.split('.')
query_clean = q.strip()
expected = hashlib.sha256(f"{ts}|{query_clean}|{self.secret}".encode()).hexdigest()
if sig != expected or (time.time() - float(ts)) > 60:
abort(403)
except: abort(403)
context_text = data.get('context', '')
if not self.api_key or not q: if not self.api_key or not q:
return Response("Error: Missing Key or Query", status=400) return Response("Error: Missing Key or Query", status=400)
@@ -141,12 +140,13 @@ class SXNGPlugin(Plugin):
context_list = [f"[{i+1}] {r.get('title')}: {r.get('content')}" for i, r in enumerate(raw_results[:6])] context_list = [f"[{i+1}] {r.get('title')}: {r.get('content')}" for i, r in enumerate(raw_results[:6])]
context_str = "\n".join(context_list) context_str = "\n".join(context_list)
# Handshake token ts = str(time.time())
tk = secrets.token_hex(16) q_clean = search.search_query.query.strip()
self.valid_tokens[tk] = time.time() + 60 sig = hashlib.sha256(f"{ts}|{q_clean}|{self.secret}".encode()).hexdigest()
tk = f"{ts}.{sig}"
b64_context = base64.b64encode(context_str.encode('utf-8')).decode('utf-8') b64_context = base64.b64encode(context_str.encode('utf-8')).decode('utf-8')
js_q = json.dumps(search.search_query.query) js_q = json.dumps(q_clean)
html_payload = f''' html_payload = f'''
<style> <style>
@@ -170,7 +170,6 @@ class SXNGPlugin(Plugin):
const tk = "{tk}"; const tk = "{tk}";
const shell = document.getElementById('sxng-stream-box'); const shell = document.getElementById('sxng-stream-box');
const data = document.getElementById('sxng-stream-data'); const data = document.getElementById('sxng-stream-data');
const loading = document.getElementById('sxng-loading');
const container = document.getElementById('urls') || document.getElementById('main_results'); const container = document.getElementById('urls') || document.getElementById('main_results');
if (container && shell) {{ container.prepend(shell); shell.style.display = 'block'; }} if (container && shell) {{ container.prepend(shell); shell.style.display = 'block'; }}
+10 -4
View File
@@ -114,10 +114,16 @@ class PluginTestCase(unittest.TestCase):
self.assertIn('/gemini-stream', content) self.assertIn('/gemini-stream', content)
def test_stream_endpoint(self): def test_stream_endpoint(self):
# Trigger index to generate a token in the plugin instance # Trigger index to generate a response containing the token
self.app.get('/') response = self.app.get('/')
# Extract the last generated token content = response.data.decode('utf-8')
token = list(plugin.valid_tokens.keys())[-1]
# Extract the token from the injected script (tk = "...")
import re
match = re.search(r'const tk = "(.*?)";', content)
if not match:
self.fail("Handshake token not found in injection")
token = match.group(1)
# Check for the appropriate key based on provider # Check for the appropriate key based on provider
key = os.getenv("OPENROUTER_API_KEY") if plugin.provider == 'openrouter' else os.getenv("GEMINI_API_KEY") key = os.getenv("OPENROUTER_API_KEY") if plugin.provider == 'openrouter' else os.getenv("GEMINI_API_KEY")