jae
e73b74cfa2
- New /api/agent/chat endpoint with Venice tool-calling loop (max 8 iter)
- Tiered models: glm-4.7-flash default, kimi-k2-thinking for Elite+
- Wallet auth: /api/auth/{nonce,verify,whoami,logout} with Ed25519 + JWT
- 10 tools registered: site search, crypto prices, SITREP, .sol lookup,
wallet xray, contraband/awesomelist search, changelog, trigger_effect
- Per-tool rate limits, 30s timeout, \$30/mo budget guard
- Frontend: tier badge, tool call cards, wallet sign-in handshake
- Changelog v1.40.0
261 lines
8.1 KiB
Python
261 lines
8.1 KiB
Python
"""Wallet-based authentication for JAE-AI agent.
|
||
|
||
Endpoints:
|
||
POST /api/auth/nonce – issue sign-in nonce for a Solana pubkey
|
||
POST /api/auth/verify – verify Ed25519 signature, issue JWT cookie
|
||
GET /api/auth/whoami – introspect cookie + holdings → tier
|
||
POST /api/auth/logout – clear cookie
|
||
"""
|
||
import json
|
||
import os
|
||
import secrets
|
||
import time
|
||
import datetime
|
||
from pathlib import Path
|
||
|
||
from flask import Blueprint, request, jsonify, make_response
|
||
import jwt
|
||
import requests as req
|
||
import base58
|
||
from nacl.signing import VerifyKey
|
||
from nacl.exceptions import BadSignatureError
|
||
|
||
from agent_tiers import compute_tier, pick_model
|
||
|
||
auth_bp = Blueprint('auth_bp', __name__, url_prefix='/api/auth')
|
||
|
||
DATA_DIR = Path(__file__).parent / 'data'
|
||
APIKEYS_FILE = DATA_DIR / 'apikeys.json'
|
||
|
||
# Solana RPC endpoints (mainnet). Helius used if key present, else public RPC.
|
||
PUBLIC_RPC = 'https://api.mainnet-beta.solana.com'
|
||
|
||
# In-memory nonce store: address -> (nonce, expires_ts)
|
||
_NONCES: dict[str, tuple[str, float]] = {}
|
||
_NONCE_TTL = 300 # 5 minutes
|
||
|
||
# Balance cache: address -> (lamports, ts)
|
||
_BALANCE_CACHE: dict[str, tuple[int, float]] = {}
|
||
_BALANCE_TTL = 60 # 1 minute
|
||
|
||
COOKIE_NAME = 'jae_session'
|
||
JWT_ALGO = 'HS256'
|
||
JWT_EXPIRY_HOURS = 24
|
||
|
||
|
||
def _load_apikeys() -> dict:
|
||
try:
|
||
with open(APIKEYS_FILE) as f:
|
||
return json.load(f)
|
||
except Exception:
|
||
return {}
|
||
|
||
|
||
def _save_apikeys(keys: dict):
|
||
try:
|
||
with open(APIKEYS_FILE, 'w') as f:
|
||
json.dump(keys, f, indent=2)
|
||
except Exception as e:
|
||
print(f'[auth] Failed to save apikeys: {e}')
|
||
|
||
|
||
def get_jwt_secret() -> str:
|
||
"""Load JAE_JWT_SECRET from apikeys.json; generate + persist if missing."""
|
||
keys = _load_apikeys()
|
||
secret = keys.get('jae_jwt_secret') or ''
|
||
if not secret:
|
||
secret = secrets.token_hex(32)
|
||
keys['jae_jwt_secret'] = secret
|
||
_save_apikeys(keys)
|
||
return secret
|
||
|
||
|
||
def get_rpc_url() -> str:
|
||
keys = _load_apikeys()
|
||
# Check custom slots for Helius
|
||
for slot in ('custom1', 'custom2', 'custom3'):
|
||
c = keys.get(slot) or {}
|
||
name = (c.get('name') or '').lower()
|
||
if 'helius' in name and c.get('key'):
|
||
return f"https://mainnet.helius-rpc.com/?api-key={c['key']}"
|
||
# Check a direct helius field
|
||
h = keys.get('helius') or {}
|
||
if h.get('api_key'):
|
||
return f"https://mainnet.helius-rpc.com/?api-key={h['api_key']}"
|
||
return PUBLIC_RPC
|
||
|
||
|
||
def fetch_sol_balance(address: str) -> float:
|
||
"""Return SOL balance (float). Cached 60s. Returns 0.0 on error."""
|
||
now = time.time()
|
||
cached = _BALANCE_CACHE.get(address)
|
||
if cached and now - cached[1] < _BALANCE_TTL:
|
||
return cached[0] / 1_000_000_000
|
||
try:
|
||
r = req.post(
|
||
get_rpc_url(),
|
||
json={'jsonrpc': '2.0', 'id': 1, 'method': 'getBalance', 'params': [address]},
|
||
timeout=6,
|
||
)
|
||
r.raise_for_status()
|
||
lamports = int(r.json()['result']['value'])
|
||
_BALANCE_CACHE[address] = (lamports, now)
|
||
return lamports / 1_000_000_000
|
||
except Exception as e:
|
||
print(f'[auth] getBalance failed for {address[:8]}…: {e}')
|
||
return 0.0
|
||
|
||
|
||
def _cleanup_nonces():
|
||
now = time.time()
|
||
stale = [k for k, (_, exp) in _NONCES.items() if exp < now]
|
||
for k in stale:
|
||
_NONCES.pop(k, None)
|
||
|
||
|
||
def _valid_solana_address(addr: str) -> bool:
|
||
if not isinstance(addr, str) or not (32 <= len(addr) <= 44):
|
||
return False
|
||
try:
|
||
raw = base58.b58decode(addr)
|
||
return len(raw) == 32
|
||
except Exception:
|
||
return False
|
||
|
||
|
||
def read_session() -> dict | None:
|
||
"""Decode JWT cookie. Returns payload or None."""
|
||
token = request.cookies.get(COOKIE_NAME)
|
||
if not token:
|
||
return None
|
||
try:
|
||
return jwt.decode(token, get_jwt_secret(), algorithms=[JWT_ALGO])
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
# ── Endpoints ────────────────────────────────────────────────────────────
|
||
|
||
@auth_bp.route('/nonce', methods=['POST'])
|
||
def issue_nonce():
|
||
data = request.get_json(silent=True) or {}
|
||
address = (data.get('address') or '').strip()
|
||
if not _valid_solana_address(address):
|
||
return jsonify({'error': 'Invalid Solana address'}), 400
|
||
_cleanup_nonces()
|
||
nonce = secrets.token_hex(16)
|
||
expires_ts = time.time() + _NONCE_TTL
|
||
expires_iso = datetime.datetime.utcfromtimestamp(expires_ts).isoformat() + 'Z'
|
||
_NONCES[address] = (nonce, expires_ts)
|
||
message = f'Sign in to jaeswift.xyz\nnonce: {nonce}\nexpires: {expires_iso}'
|
||
return jsonify({'message': message, 'nonce': nonce, 'expires': expires_iso})
|
||
|
||
|
||
@auth_bp.route('/verify', methods=['POST'])
|
||
def verify_signature():
|
||
data = request.get_json(silent=True) or {}
|
||
address = (data.get('address') or '').strip()
|
||
signature_b58 = (data.get('signature') or '').strip()
|
||
nonce = (data.get('nonce') or '').strip()
|
||
|
||
if not _valid_solana_address(address):
|
||
return jsonify({'error': 'Invalid address'}), 400
|
||
if not signature_b58 or not nonce:
|
||
return jsonify({'error': 'Missing signature or nonce'}), 400
|
||
|
||
stored = _NONCES.get(address)
|
||
if not stored:
|
||
return jsonify({'error': 'No active nonce for this address'}), 401
|
||
stored_nonce, expires_ts = stored
|
||
if stored_nonce != nonce:
|
||
return jsonify({'error': 'Nonce mismatch'}), 401
|
||
if time.time() > expires_ts:
|
||
_NONCES.pop(address, None)
|
||
return jsonify({'error': 'Nonce expired'}), 401
|
||
|
||
# Rebuild exact signed message
|
||
expires_iso = datetime.datetime.utcfromtimestamp(expires_ts).isoformat() + 'Z'
|
||
message = f'Sign in to jaeswift.xyz\nnonce: {nonce}\nexpires: {expires_iso}'
|
||
|
||
try:
|
||
pubkey_bytes = base58.b58decode(address)
|
||
vk = VerifyKey(pubkey_bytes)
|
||
# Try base58 then base64 decode for signature
|
||
sig_bytes = None
|
||
try:
|
||
sig_bytes = base58.b58decode(signature_b58)
|
||
except Exception:
|
||
import base64
|
||
sig_bytes = base64.b64decode(signature_b58)
|
||
vk.verify(message.encode('utf-8'), sig_bytes)
|
||
except BadSignatureError:
|
||
return jsonify({'error': 'Signature verification failed'}), 401
|
||
except Exception as e:
|
||
return jsonify({'error': f'Signature decode error: {e}'}), 400
|
||
|
||
# Consume nonce
|
||
_NONCES.pop(address, None)
|
||
|
||
# Compute tier based on live SOL balance
|
||
balance = fetch_sol_balance(address)
|
||
tier = compute_tier(address, balance)
|
||
|
||
payload = {
|
||
'address': address,
|
||
'tier': tier,
|
||
'iat': int(time.time()),
|
||
'exp': int(time.time()) + JWT_EXPIRY_HOURS * 3600,
|
||
}
|
||
token = jwt.encode(payload, get_jwt_secret(), algorithm=JWT_ALGO)
|
||
|
||
resp = make_response(jsonify({
|
||
'authenticated': True,
|
||
'address': address,
|
||
'tier': tier,
|
||
'balance_sol': round(balance, 4),
|
||
'model': pick_model(tier),
|
||
'expires_in': JWT_EXPIRY_HOURS * 3600,
|
||
}))
|
||
resp.set_cookie(
|
||
COOKIE_NAME, token,
|
||
max_age=JWT_EXPIRY_HOURS * 3600,
|
||
httponly=True,
|
||
secure=True,
|
||
samesite='Lax',
|
||
path='/',
|
||
)
|
||
return resp
|
||
|
||
|
||
@auth_bp.route('/whoami', methods=['GET'])
|
||
def whoami():
|
||
sess = read_session()
|
||
if not sess:
|
||
return jsonify({
|
||
'authenticated': False,
|
||
'address': None,
|
||
'tier': 'anonymous',
|
||
'model': pick_model('anonymous'),
|
||
'balance_sol': 0,
|
||
'holdings': {},
|
||
})
|
||
address = sess.get('address')
|
||
# Re-fetch balance (cached) to refresh tier on the fly
|
||
balance = fetch_sol_balance(address) if address else 0.0
|
||
tier = compute_tier(address, balance)
|
||
return jsonify({
|
||
'authenticated': True,
|
||
'address': address,
|
||
'tier': tier,
|
||
'model': pick_model(tier),
|
||
'balance_sol': round(balance, 4),
|
||
'holdings': {'sol': round(balance, 4)},
|
||
'expires_at': sess.get('exp'),
|
||
})
|
||
|
||
|
||
@auth_bp.route('/logout', methods=['POST'])
|
||
def logout():
|
||
resp = make_response(jsonify({'ok': True}))
|
||
resp.set_cookie(COOKIE_NAME, '', expires=0, path='/', secure=True, httponly=True, samesite='Lax')
|
||
return resp
|