Compare commits

..

1 Commits

Author SHA1 Message Date
ي
afcb3d5478 Add token encryption service and Wix token rotation support 2026-05-18 15:57:47 +05:30
3 changed files with 211 additions and 263 deletions

View File

@@ -9,26 +9,27 @@ from typing import Optional, Dict, Any, List
from datetime import datetime, timedelta
from loguru import logger
from services.database import get_user_db_path
from services.token_crypto_service import TokenCryptoService
class WixOAuthService:
"""Manages Wix OAuth2 authentication flow and token storage."""
def __init__(self, db_path: Optional[str] = None):
self.db_path = db_path
self.token_crypto = TokenCryptoService()
def _get_db_path(self, user_id: str) -> str:
if self.db_path:
return self.db_path
return get_user_db_path(user_id)
def _init_db(self, user_id: str):
"""Initialize database tables for OAuth tokens."""
db_path = self._get_db_path(user_id)
# Ensure directory exists
os.makedirs(os.path.dirname(db_path), exist_ok=True)
with sqlite3.connect(db_path) as conn:
cursor = conn.cursor()
cursor.execute('''
@@ -45,168 +46,133 @@ class WixOAuthService:
member_id TEXT,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
is_active BOOLEAN DEFAULT TRUE
is_active BOOLEAN DEFAULT TRUE,
token_key_version TEXT,
token_key_reference TEXT
)
''')
for column_name, column_def in [
("token_key_version", "TEXT"),
("token_key_reference", "TEXT"),
]:
try:
cursor.execute(f"ALTER TABLE wix_oauth_tokens ADD COLUMN {column_name} {column_def}")
except sqlite3.OperationalError:
pass
conn.commit()
def store_tokens(
self,
user_id: str,
access_token: str,
refresh_token: Optional[str] = None,
expires_in: Optional[int] = None,
token_type: str = 'bearer',
scope: Optional[str] = None,
site_id: Optional[str] = None,
member_id: Optional[str] = None
) -> bool:
"""
Store Wix OAuth tokens in the database.
Args:
user_id: User ID (Clerk string)
access_token: Access token from Wix
refresh_token: Optional refresh token
expires_in: Optional expiration time in seconds
token_type: Token type (default: 'bearer')
scope: Optional OAuth scope
site_id: Optional Wix site ID
member_id: Optional Wix member ID
Returns:
True if tokens were stored successfully
"""
def store_tokens(self, user_id: str, access_token: str, refresh_token: Optional[str] = None,
expires_in: Optional[int] = None, token_type: str = 'bearer', scope: Optional[str] = None,
site_id: Optional[str] = None, member_id: Optional[str] = None) -> bool:
try:
# Ensure DB is initialized for this user
self._init_db(user_id)
db_path = self._get_db_path(user_id)
expires_at = None
if expires_in:
expires_at = datetime.now() + timedelta(seconds=expires_in)
expires_at = datetime.now() + timedelta(seconds=expires_in) if expires_in else None
encrypted_access_token, encrypted_refresh_token = self.token_crypto.encrypt_pair(access_token, refresh_token)
with sqlite3.connect(db_path) as conn:
cursor = conn.cursor()
cursor.execute('''
INSERT INTO wix_oauth_tokens
(user_id, access_token, refresh_token, token_type, expires_at, expires_in, scope, site_id, member_id)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
''', (user_id, access_token, refresh_token, token_type, expires_at, expires_in, scope, site_id, member_id))
INSERT INTO wix_oauth_tokens
(user_id, access_token, refresh_token, token_type, expires_at, expires_in, scope, site_id, member_id, token_key_version, token_key_reference)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
''', (
user_id,
encrypted_access_token,
encrypted_refresh_token,
token_type,
expires_at,
expires_in,
scope,
site_id,
member_id,
self.token_crypto.key_version,
self.token_crypto.key_reference,
))
conn.commit()
logger.info(f"Wix OAuth: Token inserted into database for user {user_id}")
logger.info(f"Wix OAuth: Encrypted token stored for user {user_id}")
return True
except Exception as e:
logger.error(f"Error storing Wix tokens for user {user_id}: {e}")
return False
def get_user_tokens(self, user_id: str) -> List[Dict[str, Any]]:
"""Get all active Wix tokens for a user."""
"""Get all active Wix token rows (encrypted values)."""
try:
# Ensure database tables exist to prevent 'no such table' errors
self._init_db(user_id)
db_path = self._get_db_path(user_id)
if not os.path.exists(db_path):
return []
with sqlite3.connect(db_path) as conn:
cursor = conn.cursor()
cursor.execute('''
SELECT id, access_token, refresh_token, token_type, expires_at, expires_in, scope, site_id, member_id, created_at
SELECT id, access_token, refresh_token, token_type, expires_at, expires_in, scope, site_id, member_id, created_at, token_key_version, token_key_reference
FROM wix_oauth_tokens
WHERE user_id = ? AND is_active = TRUE AND (expires_at IS NULL OR expires_at > datetime('now'))
ORDER BY created_at DESC
''', (user_id,))
tokens = []
for row in cursor.fetchall():
tokens.append({
"id": row[0],
"access_token": row[1],
"refresh_token": row[2],
"token_type": row[3],
"expires_at": row[4],
"expires_in": row[5],
"scope": row[6],
"site_id": row[7],
"member_id": row[8],
"created_at": row[9]
})
return tokens
return [{
"id": row[0], "access_token": row[1], "refresh_token": row[2], "token_type": row[3],
"expires_at": row[4], "expires_in": row[5], "scope": row[6], "site_id": row[7],
"member_id": row[8], "created_at": row[9], "token_key_version": row[10],
"token_key_reference": row[11]
} for row in cursor.fetchall()]
except Exception as e:
logger.error(f"Error getting Wix tokens for user {user_id}: {e}")
return []
def get_user_token_status(self, user_id: str) -> Dict[str, Any]:
"""Get detailed token status for a user including expired tokens."""
try:
# Ensure database tables exist to prevent 'no such table' errors
self._init_db(user_id)
def get_user_tokens_decrypted(self, user_id: str) -> List[Dict[str, Any]]:
"""Decrypt tokens for integration managers and token refresh routines."""
decrypted = []
for token in self.get_user_tokens(user_id):
token_copy = dict(token)
token_copy["access_token"] = self.token_crypto.decrypt_token(token_copy.get("access_token"))
token_copy["refresh_token"] = self.token_crypto.decrypt_token(token_copy.get("refresh_token"))
decrypted.append(token_copy)
return decrypted
def get_user_token_status(self, user_id: str) -> Dict[str, Any]:
try:
self._init_db(user_id)
db_path = self._get_db_path(user_id)
if not os.path.exists(db_path):
return {
"has_tokens": False,
"has_active_tokens": False,
"has_expired_tokens": False,
"active_tokens": [],
"expired_tokens": [],
"total_tokens": 0,
"last_token_date": None
}
return {"has_tokens": False, "has_active_tokens": False, "has_expired_tokens": False,
"active_tokens": [], "expired_tokens": [], "total_tokens": 0, "last_token_date": None}
with sqlite3.connect(db_path) as conn:
cursor = conn.cursor()
# Get all tokens (active and expired)
cursor.execute('''
SELECT id, access_token, refresh_token, token_type, expires_at, expires_in, scope, site_id, member_id, created_at, is_active
SELECT id, access_token, refresh_token, token_type, expires_at, expires_in, scope, site_id, member_id,
created_at, is_active, token_key_version, token_key_reference
FROM wix_oauth_tokens
WHERE user_id = ?
ORDER BY created_at DESC
''', (user_id,))
all_tokens = []
active_tokens = []
expired_tokens = []
all_tokens, active_tokens, expired_tokens = [], [], []
for row in cursor.fetchall():
token_data = {
"id": row[0],
"access_token": row[1],
"refresh_token": row[2],
"token_type": row[3],
"expires_at": row[4],
"expires_in": row[5],
"scope": row[6],
"site_id": row[7],
"member_id": row[8],
"created_at": row[9],
"is_active": bool(row[10])
"id": row[0], "access_token": row[1], "refresh_token": row[2], "token_type": row[3],
"expires_at": row[4], "expires_in": row[5], "scope": row[6], "site_id": row[7],
"member_id": row[8], "created_at": row[9], "is_active": bool(row[10]),
"token_key_version": row[11], "token_key_reference": row[12]
}
all_tokens.append(token_data)
# Determine expiry using robust parsing and is_active flag
is_active_flag = bool(row[10])
not_expired = False
try:
expires_at_val = row[4]
if expires_at_val:
# First try Python parsing
try:
dt = datetime.fromisoformat(expires_at_val) if isinstance(expires_at_val, str) else expires_at_val
not_expired = dt > datetime.now()
except Exception:
# Fallback to SQLite comparison
cursor.execute("SELECT datetime('now') < ?", (expires_at_val,))
not_expired = cursor.fetchone()[0] == 1
else:
# No expiry stored => consider not expired
not_expired = True
except Exception:
not_expired = False
@@ -215,7 +181,7 @@ class WixOAuthService:
active_tokens.append(token_data)
else:
expired_tokens.append(token_data)
return {
"has_tokens": len(all_tokens) > 0,
"has_active_tokens": len(active_tokens) > 0,
@@ -225,81 +191,101 @@ class WixOAuthService:
"total_tokens": len(all_tokens),
"last_token_date": all_tokens[0]["created_at"] if all_tokens else None
}
except Exception as e:
logger.error(f"Error getting Wix token status for user {user_id}: {e}")
return {
"has_tokens": False,
"has_active_tokens": False,
"has_expired_tokens": False,
"active_tokens": [],
"expired_tokens": [],
"total_tokens": 0,
"last_token_date": None,
"error": str(e)
}
def update_tokens(
self,
user_id: str,
access_token: str,
refresh_token: Optional[str] = None,
expires_in: Optional[int] = None
) -> bool:
"""Update tokens for a user (e.g., after refresh)."""
return {"has_tokens": False, "has_active_tokens": False, "has_expired_tokens": False,
"active_tokens": [], "expired_tokens": [], "total_tokens": 0, "last_token_date": None, "error": str(e)}
def update_tokens(self, user_id: str, access_token: str, refresh_token: Optional[str] = None,
expires_in: Optional[int] = None) -> bool:
try:
# Ensure DB initialized for this user
self._init_db(user_id)
db_path = self._get_db_path(user_id)
expires_at = datetime.now() + timedelta(seconds=expires_in) if expires_in else None
encrypted_access_token = self.token_crypto.encrypt_token(access_token)
encrypted_refresh_token = self.token_crypto.encrypt_token(refresh_token) if refresh_token else None
expires_at = None
if expires_in:
expires_at = datetime.now() + timedelta(seconds=expires_in)
with sqlite3.connect(db_path) as conn:
cursor = conn.cursor()
if refresh_token:
cursor.execute('''
UPDATE wix_oauth_tokens
SET access_token = ?, refresh_token = ?, expires_at = ?, expires_in = ?,
is_active = TRUE, updated_at = datetime('now')
WHERE user_id = ? AND refresh_token = ?
''', (access_token, refresh_token, expires_at, expires_in, user_id, refresh_token))
UPDATE wix_oauth_tokens
SET access_token = ?, refresh_token = ?, expires_at = ?, expires_in = ?,
is_active = TRUE, updated_at = datetime('now'), token_key_version = ?, token_key_reference = ?
WHERE user_id = ? AND (refresh_token = ? OR refresh_token = ?)
''', (encrypted_access_token, encrypted_refresh_token, expires_at, expires_in,
self.token_crypto.key_version, self.token_crypto.key_reference,
user_id, encrypted_refresh_token, refresh_token))
else:
cursor.execute('''
UPDATE wix_oauth_tokens
SET access_token = ?, expires_at = ?, expires_in = ?,
is_active = TRUE, updated_at = datetime('now')
UPDATE wix_oauth_tokens
SET access_token = ?, expires_at = ?, expires_in = ?,
is_active = TRUE, updated_at = datetime('now'), token_key_version = ?, token_key_reference = ?
WHERE user_id = ? AND id = (SELECT id FROM wix_oauth_tokens WHERE user_id = ? ORDER BY created_at DESC LIMIT 1)
''', (access_token, expires_at, expires_in, user_id, user_id))
''', (encrypted_access_token, expires_at, expires_in,
self.token_crypto.key_version, self.token_crypto.key_reference, user_id, user_id))
conn.commit()
logger.info(f"Wix OAuth: Tokens updated for user {user_id}")
logger.info(f"Wix OAuth: Encrypted tokens updated for user {user_id}")
return True
except Exception as e:
logger.error(f"Error updating Wix tokens for user {user_id}: {e}")
return False
def rotate_token_encryption(self, user_id: str, batch_size: int = 100) -> Dict[str, int]:
"""Re-encrypt existing token rows in batches for key rotation."""
self._init_db(user_id)
db_path = self._get_db_path(user_id)
rotated, skipped, last_id = 0, 0, 0
with sqlite3.connect(db_path) as conn:
cursor = conn.cursor()
while True:
cursor.execute('''
SELECT id, access_token, refresh_token
FROM wix_oauth_tokens
WHERE user_id = ? AND id > ?
ORDER BY id ASC
LIMIT ?
''', (user_id, last_id, batch_size))
rows = cursor.fetchall()
if not rows:
break
for row_id, enc_access, enc_refresh in rows:
last_id = row_id
try:
plain_access = self.token_crypto.decrypt_token(enc_access)
plain_refresh = self.token_crypto.decrypt_token(enc_refresh) if enc_refresh else None
except Exception:
skipped += 1
continue
new_access, new_refresh = self.token_crypto.encrypt_pair(plain_access, plain_refresh)
cursor.execute('''
UPDATE wix_oauth_tokens
SET access_token = ?, refresh_token = ?, token_key_version = ?, token_key_reference = ?, updated_at = datetime('now')
WHERE id = ?
''', (new_access, new_refresh, self.token_crypto.key_version, self.token_crypto.key_reference, row_id))
rotated += 1
conn.commit()
logger.info(f"Wix OAuth: Encryption rotation complete for user {user_id}; rotated={rotated}, skipped={skipped}")
return {"rotated": rotated, "skipped": skipped}
def revoke_token(self, user_id: str, token_id: int) -> bool:
"""Revoke a Wix OAuth token."""
try:
db_path = self._get_db_path(user_id)
with sqlite3.connect(db_path) as conn:
cursor = conn.cursor()
cursor.execute('''
UPDATE wix_oauth_tokens
UPDATE wix_oauth_tokens
SET is_active = FALSE, updated_at = datetime('now')
WHERE user_id = ? AND id = ?
''', (user_id, token_id))
conn.commit()
if cursor.rowcount > 0:
logger.info(f"Wix token {token_id} revoked for user {user_id}")
return True
return False
except Exception as e:
logger.error(f"Error revoking Wix token: {e}")
return False

View File

@@ -101,7 +101,6 @@ class AgentContextVFS:
"/steps/integrations": AgentFlatContextStore.STEP5_FILENAME,
}
HIGH_SIGNAL_MARKERS = ("agent_summary", "high_signal_terms", "quick_facts", "context_type")
LOW_CONFIDENCE_MARKER = "low_confidence"
def __init__(self, user_id: str, project_id: Optional[str] = None):
self.user_id = user_id
@@ -295,101 +294,6 @@ class AgentContextVFS:
)
return ranked[: max(1, top_k)]
@staticmethod
def _mnemonic_token(result: Dict[str, Any], rank: int) -> str:
"""Create compressed mnemonic token with source reference."""
path = str(result.get("path") or "unknown")
reason = str(result.get("reason") or "match")
confidence = float(result.get("confidence") or 0.0)
low_flag = "!" if result.get(AgentContextVFS.LOW_CONFIDENCE_MARKER) else ""
src = path.replace(".json", "").replace("_", "-")[:28]
hint = reason.replace(" ", "-")[:20]
return f"M{rank}:{src}|{hint}|c{confidence:.2f}{low_flag}"
@staticmethod
def _detect_contradictions(results: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
"""Detect contradictory learnings by path with conflicting reasons/relevance classes."""
by_path: Dict[str, List[Dict[str, Any]]] = {}
for item in results:
p = str(item.get("path") or "")
by_path.setdefault(p, []).append(item)
contradictions: List[Dict[str, Any]] = []
for path, rows in by_path.items():
reasons = {str(r.get("reason") or "").strip().lower() for r in rows}
relevance = {str(r.get("relevance") or "").strip().lower() for r in rows}
# contradictory if both high/supported or mixed summary/body signals in same source cluster
if len(reasons) > 1 and len(relevance) > 1:
contradictions.append(
{
"path": path,
"reason_variants": sorted([r for r in reasons if r]),
"relevance_variants": sorted([r for r in relevance if r]),
"count": len(rows),
}
)
return contradictions
def _run_synthesis_pipeline(
self, ranked_results: List[Dict[str, Any]], *, char_budget: int = 1200, top_k: int = 5
) -> Dict[str, Any]:
"""
Flat-context synthesis pipeline:
1) Compress telemetry into mnemonic tokens with source references
2) Detect contradictions and mark low-confidence heuristics
3) Select top-ranked, budget-fitting tokens for prompt injection
4) Persist synthesis + source lineage for explainability
"""
contradictions = self._detect_contradictions(ranked_results)
contradiction_paths = {c["path"] for c in contradictions}
normalized: List[Dict[str, Any]] = []
for idx, item in enumerate(ranked_results, start=1):
row = dict(item)
low_conf = bool(row.get("low_probability")) or (str(row.get("path") or "") in contradiction_paths)
row[self.LOW_CONFIDENCE_MARKER] = low_conf
if low_conf:
row["confidence"] = round(max(0.05, float(row.get("confidence", 0.0)) * 0.7), 3)
row["mnemonic_token"] = self._mnemonic_token(row, idx)
normalized.append(row)
chosen: List[Dict[str, Any]] = []
used = 0
for row in normalized[: max(1, top_k * 3)]:
token = str(row.get("mnemonic_token") or "")
cost = len(token) + 8
if chosen and used + cost > char_budget:
continue
chosen.append(row)
used += cost
if len(chosen) >= top_k:
break
synthesis = {
"created_at": datetime.now(timezone.utc).isoformat(),
"top_k": top_k,
"char_budget": char_budget,
"char_budget_used": used,
"selected_mnemonics": [c.get("mnemonic_token") for c in chosen],
"source_lineage": [
{
"mnemonic_token": c.get("mnemonic_token"),
"path": c.get("path"),
"reason": c.get("reason"),
"confidence": c.get("confidence"),
"low_confidence": c.get(self.LOW_CONFIDENCE_MARKER, False),
}
for c in chosen
],
"contradictions": contradictions,
}
self.append_activity_log(
event_type="flat_context_synthesis",
actor="agent_context_vfs",
details=synthesis,
)
return {"ranked_results": normalized, "synthesis": synthesis}
@staticmethod
def _resolve_json_path(data: Any, path_query: str) -> Any:
"""Resolve dot/bracket JSON path such as 'data.seo_audit.recommendations[0]'."""
@@ -614,26 +518,15 @@ class AgentContextVFS:
bounded_results.append(r)
used += cost
synthesis_bundle = self._run_synthesis_pipeline(
self._static_triage(bounded_results, normalized),
char_budget=1200,
top_k=5,
)
triaged_results = synthesis_bundle["ranked_results"]
synthesis = synthesis_bundle["synthesis"]
result = {
"query": normalized,
"attempted_queries": attempted_queries,
"matched_files_count": len(matched_files),
"results": triaged_results,
"results": self._static_triage(bounded_results, normalized),
"notice": notice,
"char_budget_used": used,
"can_answer": bool(bounded_results),
"synthesis": synthesis,
"prompt_context_mnemonics": synthesis.get("selected_mnemonics", []),
}
# Top-ranked, budget-fitting mnemonic tokens are the only ones intended for prompt context injection.
result["triage_top5"] = self._llm_router_stub(result["results"], top_k=5)
logger.info(
f"[vfs_audit] user={self.store.safe_user_id} action=search_context query={normalized!r} results={len(result['results'])}"

View File

@@ -0,0 +1,69 @@
"""Service for encrypting/decrypting integration tokens with key version metadata."""
import base64
import hashlib
import os
from typing import Optional, Tuple
from cryptography.fernet import Fernet, InvalidToken
from loguru import logger
class TokenCryptoService:
"""Token encryption/decryption service with key version support."""
ENV_KEY = "ALWRITY_TOKEN_ENCRYPTION_KEY"
ENV_KEY_VERSION = "ALWRITY_TOKEN_KEY_VERSION"
def __init__(self):
raw_key = os.getenv(self.ENV_KEY, "")
if raw_key:
self._fernet_key = self._normalize_key(raw_key)
else:
self._fernet_key = self._derive_dev_key()
self._fernet = Fernet(self._fernet_key)
self._key_version = os.getenv(self.ENV_KEY_VERSION, "v1")
self._key_reference = self._fingerprint(self._fernet_key)
@property
def key_version(self) -> str:
return self._key_version
@property
def key_reference(self) -> str:
return self._key_reference
def encrypt_token(self, token: Optional[str]) -> Optional[str]:
if token is None:
return None
return self._fernet.encrypt(token.encode("utf-8")).decode("utf-8")
def decrypt_token(self, encrypted_token: Optional[str]) -> Optional[str]:
if encrypted_token is None:
return None
try:
return self._fernet.decrypt(encrypted_token.encode("utf-8")).decode("utf-8")
except InvalidToken:
logger.error("Token decryption failed due to invalid token/key")
raise
def encrypt_pair(self, access_token: str, refresh_token: Optional[str]) -> Tuple[str, Optional[str]]:
return self.encrypt_token(access_token), self.encrypt_token(refresh_token)
@staticmethod
def _normalize_key(raw_key: str) -> bytes:
raw_key = raw_key.strip()
if len(raw_key) == 44 and raw_key.endswith("="):
return raw_key.encode("utf-8")
digest = hashlib.sha256(raw_key.encode("utf-8")).digest()
return base64.urlsafe_b64encode(digest)
@staticmethod
def _derive_dev_key() -> bytes:
seed = "alwrity-local-token-key"
digest = hashlib.sha256(seed.encode("utf-8")).digest()
return base64.urlsafe_b64encode(digest)
@staticmethod
def _fingerprint(key: bytes) -> str:
return hashlib.sha256(key).hexdigest()[:16]