diff --git a/.gitignore b/.gitignore index 54c2d1b842..d9dca2a054 100644 --- a/.gitignore +++ b/.gitignore @@ -5,3 +5,26 @@ __pycache__/ *.pyc .pytest_cache/ mempal.yaml + +# Security-sensitive files +.env +.env.* +*.key +auth_token +entities.json + +# Database files +*.sqlite3 + +# Virtual environments +venv/ +.venv/ +env/ + +# IDE +.idea/ +.vscode/ + +# Test coverage +htmlcov/ +.coverage diff --git a/mempalace/config.py b/mempalace/config.py index 5a736509f9..ecdfdfec25 100644 --- a/mempalace/config.py +++ b/mempalace/config.py @@ -8,6 +8,8 @@ import os from pathlib import Path +from .security import secure_dir, secure_file + DEFAULT_PALACE_PATH = os.path.expanduser("~/.mempalace/palace") DEFAULT_COLLECTION_NAME = "mempalace_drawers" @@ -123,9 +125,39 @@ def hall_keywords(self): """Mapping of hall names to keyword lists.""" return self._file_config.get("hall_keywords", DEFAULT_HALL_KEYWORDS) + # ── Security settings ──────────────────────────────────────────────── + + def _security(self, key, default=None): + """Read a value from the nested security config block.""" + env_key = f"MEMPALACE_{key.upper()}" + env_val = os.environ.get(env_key) + if env_val is not None: + if isinstance(default, bool): + return env_val.lower() in ("1", "true", "yes") + if isinstance(default, int): + return int(env_val) + return env_val + return self._file_config.get("security", {}).get(key, default) + + @property + def auth_enabled(self): + """Whether MCP token authentication is required.""" + return self._security("auth_enabled", False) + + @property + def encryption_enabled(self): + """Whether data-at-rest encryption is active.""" + return self._security("encryption_enabled", False) + + @property + def max_content_size(self): + """Maximum allowed content size in bytes for MCP write operations.""" + return self._security("max_content_size", 1_048_576) # 1 MB + def init(self): """Create config directory and write default config.json if it doesn't exist.""" self._config_dir.mkdir(parents=True, exist_ok=True) + secure_dir(self._config_dir) if not self._config_file.exists(): default_config = { "palace_path": DEFAULT_PALACE_PATH, @@ -135,6 +167,7 @@ def init(self): } with open(self._config_file, "w") as f: json.dump(default_config, f, indent=2) + secure_file(self._config_file) return self._config_file def save_people_map(self, people_map): @@ -146,4 +179,5 @@ def save_people_map(self, people_map): self._config_dir.mkdir(parents=True, exist_ok=True) with open(self._people_map_file, "w") as f: json.dump(people_map, f, indent=2) + secure_file(self._people_map_file) return self._people_map_file diff --git a/mempalace/convo_miner.py b/mempalace/convo_miner.py index c316407fdb..662d370442 100644 --- a/mempalace/convo_miner.py +++ b/mempalace/convo_miner.py @@ -10,7 +10,7 @@ import os import sys -import hashlib +from .security import content_hash from pathlib import Path from datetime import datetime from collections import defaultdict @@ -356,7 +356,7 @@ def mine_convos( chunk_room = chunk.get("memory_type", room) if extract_mode == "general" else room if extract_mode == "general": room_counts[chunk_room] += 1 - drawer_id = f"drawer_{wing}_{chunk_room}_{hashlib.md5((source_file + str(chunk['chunk_index'])).encode(), usedforsecurity=False).hexdigest()[:16]}" + drawer_id = f"drawer_{wing}_{chunk_room}_{content_hash(source_file + str(chunk['chunk_index']))}" try: collection.add( documents=[chunk["content"]], diff --git a/mempalace/knowledge_graph.py b/mempalace/knowledge_graph.py index 226c92da54..785df7e7c7 100644 --- a/mempalace/knowledge_graph.py +++ b/mempalace/knowledge_graph.py @@ -35,20 +35,22 @@ kg.invalidate("Max", "has_issue", "sports_injury", ended="2026-02-15") """ -import hashlib import json import os import sqlite3 from datetime import date, datetime from pathlib import Path +from .security import content_hash, decrypt as sec_decrypt, encrypt as sec_encrypt, secure_file + DEFAULT_KG_PATH = os.path.expanduser("~/.mempalace/knowledge_graph.sqlite3") class KnowledgeGraph: - def __init__(self, db_path: str = None): + def __init__(self, db_path: str = None, fernet=None): self.db_path = db_path or DEFAULT_KG_PATH + self._fernet = fernet Path(self.db_path).parent.mkdir(parents=True, exist_ok=True) self._init_db() @@ -85,6 +87,7 @@ def _init_db(self): """) conn.commit() conn.close() + secure_file(self.db_path) def _conn(self): conn = sqlite3.connect(self.db_path, timeout=10) @@ -96,10 +99,36 @@ def _entity_id(self, name: str) -> str: # ── Write operations ────────────────────────────────────────────────── + def _encrypt_props(self, props_json: str) -> str: + """Encrypt properties JSON if encryption is enabled.""" + if self._fernet: + return sec_encrypt(self._fernet, props_json) + return props_json + + def _decrypt_props(self, props_str: str) -> str: + """Decrypt properties string if it looks like Fernet ciphertext. + + Fernet tokens always start with 'gAAAAA'. If the string starts with '{' + it's unencrypted JSON from before encryption was enabled. + """ + if not props_str or props_str.startswith("{"): + return props_str # Already plaintext JSON + if self._fernet: + try: + return sec_decrypt(self._fernet, props_str) + except Exception: + import logging + + logging.getLogger("mempalace_security").error( + "Failed to decrypt entity properties — wrong key or corrupted data" + ) + return "{}" # Return empty properties rather than ciphertext + return props_str # No fernet available, return as-is + def add_entity(self, name: str, entity_type: str = "unknown", properties: dict = None): """Add or update an entity node.""" eid = self._entity_id(name) - props = json.dumps(properties or {}) + props = self._encrypt_props(json.dumps(properties or {})) conn = self._conn() conn.execute( "INSERT OR REPLACE INTO entities (id, name, type, properties) VALUES (?, ?, ?, ?)", @@ -147,7 +176,7 @@ def add_triple( conn.close() return existing[0] # Already exists and still valid - triple_id = f"t_{sub_id}_{pred}_{obj_id}_{hashlib.md5(f'{valid_from}{datetime.now().isoformat()}'.encode()).hexdigest()[:8]}" + triple_id = f"t_{sub_id}_{pred}_{obj_id}_{content_hash(f'{valid_from}{datetime.now().isoformat()}', length=8)}" conn.execute( """INSERT INTO triples (id, subject, predicate, object, valid_from, valid_to, confidence, source_closet, source_file) diff --git a/mempalace/mcp_server.py b/mempalace/mcp_server.py index 2169255525..56c036480e 100644 --- a/mempalace/mcp_server.py +++ b/mempalace/mcp_server.py @@ -20,12 +20,20 @@ import sys import json import logging -import hashlib from datetime import datetime from .config import MempalaceConfig from .version import __version__ from .searcher import search_memories +from .security import ( + content_hash, + decrypt, + encrypt, + load_or_create_key, + load_or_create_token, + validate_bind_address, + verify_token, +) from .palace_graph import traverse, find_tunnels, graph_stats import chromadb @@ -38,6 +46,33 @@ _config = MempalaceConfig() +# Auth state — loaded at startup if auth is enabled +_auth_token = None + +# Encryption state — loaded at startup if encryption is enabled +_fernet = None + + +def _decrypt_doc(doc, meta): + """Return decrypted content if encrypted_content exists in metadata. + + When encryption is enabled, the documents field contains only a placeholder + (content hash), and the real content lives in encrypted_content metadata. + """ + encrypted = meta.get("encrypted_content") + if not encrypted: + return doc # Unencrypted drawer — return as-is + + if not _fernet: + logger.error("Encrypted content found but no encryption key loaded") + return "[encrypted — key not available]" + + try: + return decrypt(_fernet, encrypted) + except Exception: + logger.error("Failed to decrypt content — wrong key or corrupted data") + return "[encrypted — decryption failed]" + def _get_collection(create=False): """Return the ChromaDB collection, or None on failure.""" @@ -177,6 +212,7 @@ def tool_search(query: str, limit: int = 5, wing: str = None, room: str = None): wing=wing, room=room, n_results=limit, + fernet=_fernet, ) @@ -197,7 +233,7 @@ def tool_check_duplicate(content: str, threshold: float = 0.9): similarity = round(1 - dist, 3) if similarity >= threshold: meta = results["metadatas"][0][i] - doc = results["documents"][0][i] + doc = _decrypt_doc(results["documents"][0][i], meta) duplicates.append( { "id": drawer_id, @@ -264,22 +300,30 @@ def tool_add_drawer( "matches": dup["matches"], } - drawer_id = f"drawer_{wing}_{room}_{hashlib.md5((content[:100] + datetime.now().isoformat()).encode()).hexdigest()[:16]}" + drawer_id = f"drawer_{wing}_{room}_{content_hash(content[:100] + datetime.now().isoformat())}" + + metadata = { + "wing": wing, + "room": room, + "source_file": source_file or "", + "chunk_index": 0, + "added_by": added_by, + "filed_at": datetime.now().isoformat(), + } + + # When encryption is enabled, store only a placeholder in the documents + # field (used by ChromaDB for embeddings) and the real content encrypted + # in metadata. This ensures plaintext is not persisted on disk. + doc_content = content + if _fernet: + metadata["encrypted_content"] = encrypt(_fernet, content) + doc_content = f"[encrypted:{content_hash(content, length=32)}]" try: col.add( ids=[drawer_id], - documents=[content], - metadatas=[ - { - "wing": wing, - "room": room, - "source_file": source_file or "", - "chunk_index": 0, - "added_by": added_by, - "filed_at": datetime.now().isoformat(), - } - ], + documents=[doc_content], + metadatas=[metadata], ) logger.info(f"Filed drawer: {drawer_id} → {wing}/{room}") return {"success": True, "drawer_id": drawer_id, "wing": wing, "room": room} @@ -361,24 +405,29 @@ def tool_diary_write(agent_name: str, entry: str, topic: str = "general"): return _no_palace() now = datetime.now() - entry_id = f"diary_{wing}_{now.strftime('%Y%m%d_%H%M%S')}_{hashlib.md5(entry[:50].encode()).hexdigest()[:8]}" + entry_id = f"diary_{wing}_{now.strftime('%Y%m%d_%H%M%S')}_{content_hash(entry[:50], length=8)}" + + metadata = { + "wing": wing, + "room": room, + "hall": "hall_diary", + "topic": topic, + "type": "diary_entry", + "agent": agent_name, + "filed_at": now.isoformat(), + "date": now.strftime("%Y-%m-%d"), + } + + doc_content = entry + if _fernet: + metadata["encrypted_content"] = encrypt(_fernet, entry) + doc_content = f"[encrypted:{content_hash(entry, length=32)}]" try: col.add( ids=[entry_id], - documents=[entry], - metadatas=[ - { - "wing": wing, - "room": room, - "hall": "hall_diary", - "topic": topic, - "type": "diary_entry", - "agent": agent_name, - "filed_at": now.isoformat(), - "date": now.strftime("%Y-%m-%d"), - } - ], + documents=[doc_content], + metadatas=[metadata], ) logger.info(f"Diary entry: {entry_id} → {wing}/diary/{topic}") return { @@ -420,7 +469,7 @@ def tool_diary_read(agent_name: str, last_n: int = 10): "date": meta.get("date", ""), "timestamp": meta.get("filed_at", ""), "topic": meta.get("topic", ""), - "content": doc, + "content": _decrypt_doc(doc, meta), } ) @@ -689,24 +738,44 @@ def tool_diary_read(agent_name: str, last_n: int = 10): } +def _check_auth(params, req_id): + """Check auth token if authentication is enabled. Returns error response or None.""" + if _auth_token is None: + return None # Auth not enabled + token = params.get("_meta", {}).get("auth_token", "") + if not verify_token(token, _auth_token): + return { + "jsonrpc": "2.0", + "id": req_id, + "error": {"code": -32001, "message": "Unauthorized — invalid or missing auth token"}, + } + return None + + def handle_request(request): method = request.get("method", "") params = request.get("params", {}) req_id = request.get("id") if method == "initialize": + server_info = {"name": "mempalace", "version": __version__} + if _auth_token is not None: + server_info["authRequired"] = True return { "jsonrpc": "2.0", "id": req_id, "result": { "protocolVersion": "2024-11-05", "capabilities": {"tools": {}}, - "serverInfo": {"name": "mempalace", "version": __version__}, + "serverInfo": server_info, }, } elif method == "notifications/initialized": return None elif method == "tools/list": + auth_err = _check_auth(params, req_id) + if auth_err: + return auth_err return { "jsonrpc": "2.0", "id": req_id, @@ -718,6 +787,9 @@ def handle_request(request): }, } elif method == "tools/call": + auth_err = _check_auth(params, req_id) + if auth_err: + return auth_err tool_name = params.get("name") tool_args = params.get("arguments", {}) if tool_name not in TOOLS: @@ -726,10 +798,53 @@ def handle_request(request): "id": req_id, "error": {"code": -32601, "message": f"Unknown tool: {tool_name}"}, } + + tool_def = TOOLS[tool_name] + schema = tool_def["input_schema"] + schema_props = schema.get("properties", {}) + required_fields = schema.get("required", []) + + # Validate required fields are present + missing = [f for f in required_fields if f not in tool_args] + if missing: + return { + "jsonrpc": "2.0", + "id": req_id, + "error": { + "code": -32602, + "message": f"Missing required fields: {', '.join(missing)}", + }, + } + + # Reject unknown fields + unknown = [k for k in tool_args if k not in schema_props and k != "_meta"] + if unknown: + return { + "jsonrpc": "2.0", + "id": req_id, + "error": { + "code": -32602, + "message": f"Unknown fields: {', '.join(unknown)}", + }, + } + + # Enforce content size limits on write operations + max_size = _config.max_content_size + for field in ("content", "entry"): + value = tool_args.get(field) + if isinstance(value, str) and len(value) > max_size: + return { + "jsonrpc": "2.0", + "id": req_id, + "error": { + "code": -32000, + "message": f"Content exceeds max size ({max_size} bytes)", + }, + } + # Coerce argument types based on input_schema. # MCP JSON transport may deliver integers as floats or strings; # ChromaDB and Python slicing require native int. - schema_props = TOOLS[tool_name]["input_schema"].get("properties", {}) for key, value in list(tool_args.items()): prop_schema = schema_props.get(key, {}) declared_type = prop_schema.get("type") @@ -737,8 +852,12 @@ def handle_request(request): tool_args[key] = int(value) elif declared_type == "number" and not isinstance(value, (int, float)): tool_args[key] = float(value) + + # Strip _meta before dispatching (used for auth, not tool args) + tool_args.pop("_meta", None) + try: - result = TOOLS[tool_name]["handler"](**tool_args) + result = tool_def["handler"](**tool_args) return { "jsonrpc": "2.0", "id": req_id, @@ -759,8 +878,38 @@ def handle_request(request): } +def start_http_server(host="127.0.0.1", port=8766): + """Start an HTTP transport for the MCP server. + + Only localhost addresses are allowed. This is a defensive guard — + the MCP server must never be exposed to the network. + + Raises ValueError if host is not localhost. + """ + validate_bind_address(host) + raise NotImplementedError( + "HTTP transport is not yet implemented. Use stdio transport (default)." + ) + + def main(): + global _auth_token, _fernet logger.info("MemPalace MCP Server starting...") + if _config.auth_enabled: + _auth_token = load_or_create_token(_config._config_dir) + logger.info("Authentication enabled.") + if _config.encryption_enabled: + _fernet = load_or_create_key(_config._config_dir) + logger.info("Encryption at rest enabled.") + + transport = _config._file_config.get("security", {}).get("transport", "stdio") + if transport == "http": + host = _config._file_config.get("security", {}).get("http_host", "127.0.0.1") + port = _config._file_config.get("security", {}).get("http_port", 8766) + start_http_server(host, port) + return + + # Default: stdio transport while True: try: line = sys.stdin.readline() diff --git a/mempalace/miner.py b/mempalace/miner.py index 7b4e9491ff..108e7e366f 100644 --- a/mempalace/miner.py +++ b/mempalace/miner.py @@ -9,7 +9,7 @@ import os import sys -import hashlib +from .security import content_hash import fnmatch from pathlib import Path from datetime import datetime @@ -415,7 +415,7 @@ def add_drawer( collection, wing: str, room: str, content: str, source_file: str, chunk_index: int, agent: str ): """Add one drawer to the palace.""" - drawer_id = f"drawer_{wing}_{room}_{hashlib.md5((source_file + str(chunk_index)).encode(), usedforsecurity=False).hexdigest()[:16]}" + drawer_id = f"drawer_{wing}_{room}_{content_hash(source_file + str(chunk_index))}" try: collection.add( documents=[content], diff --git a/mempalace/searcher.py b/mempalace/searcher.py index 163abd88c5..a69e023e5b 100644 --- a/mempalace/searcher.py +++ b/mempalace/searcher.py @@ -11,6 +11,8 @@ import chromadb +from .security import decrypt + logger = logging.getLogger("mempalace_mcp") @@ -91,7 +93,12 @@ def search(query: str, palace_path: str, wing: str = None, room: str = None, n_r def search_memories( - query: str, palace_path: str, wing: str = None, room: str = None, n_results: int = 5 + query: str, + palace_path: str, + wing: str = None, + room: str = None, + n_results: int = 5, + fernet=None, ) -> dict: """ Programmatic search — returns a dict instead of printing. @@ -135,9 +142,21 @@ def search_memories( hits = [] for doc, meta, dist in zip(docs, metas, dists): + encrypted = meta.get("encrypted_content") + if encrypted: + if fernet: + try: + text = decrypt(fernet, encrypted) + except Exception: + logger.error("Failed to decrypt search result — wrong key or corrupted data") + text = "[encrypted — decryption failed]" + else: + text = "[encrypted — key not available]" + else: + text = doc hits.append( { - "text": doc, + "text": text, "wing": meta.get("wing", "unknown"), "room": meta.get("room", "unknown"), "source_file": Path(meta.get("source_file", "?")).name, diff --git a/mempalace/security.py b/mempalace/security.py new file mode 100644 index 0000000000..0bb62c11f4 --- /dev/null +++ b/mempalace/security.py @@ -0,0 +1,242 @@ +""" +security.py — Centralized security primitives for MemPalace. + +Provides: + - File permission hardening (secure_file, secure_dir) + - SHA-256 content hashing (replacing MD5) + - Localhost bind address validation + - Token-based authentication (generate, store, verify) + - Fernet encryption at rest (encrypt, decrypt, key management) + +All security features are opt-in. MemPalace works without them. +""" + +import hashlib +import hmac +import logging +import os +import secrets +import sys +from pathlib import Path + +logger = logging.getLogger("mempalace_security") + +# ── File permissions ───────────────────────────────────────────────────── + +_IS_WINDOWS = sys.platform == "win32" + + +def secure_file(path): + """Set file to owner-only read/write (0o600). No-op on Windows.""" + if _IS_WINDOWS: + return + try: + os.chmod(path, 0o600) + except OSError as e: + logger.warning("Could not set permissions on %s: %s", path, e) + + +def secure_dir(path): + """Set directory to owner-only (0o700). No-op on Windows.""" + if _IS_WINDOWS: + return + try: + os.chmod(path, 0o700) + except OSError as e: + logger.warning("Could not set permissions on %s: %s", path, e) + + +# ── Hashing ────────────────────────────────────────────────────────────── + + +def content_hash(data, length=16): + """SHA-256 hash of data, truncated to `length` hex characters. + + Drop-in replacement for the old MD5-based ID generation. + """ + return hashlib.sha256(data.encode()).hexdigest()[:length] + + +# ── Localhost validation ───────────────────────────────────────────────── + +_LOCALHOST_ADDRESSES = frozenset({"127.0.0.1", "::1", "localhost"}) + + +def validate_bind_address(host): + """Raise ValueError if host is not a localhost address. + + Defensive guard for any future network transport — ensures the MCP + server can never accidentally bind to a public interface. + """ + if host not in _LOCALHOST_ADDRESSES: + raise ValueError( + f"Refusing to bind to non-localhost address '{host}'. " + f"Allowed: {', '.join(sorted(_LOCALHOST_ADDRESSES))}" + ) + + +# ── Authentication ─────────────────────────────────────────────────────── + +_KEYRING_SERVICE = "mempalace" + + +def generate_auth_token(): + """Generate a URL-safe random auth token (43 characters).""" + return secrets.token_urlsafe(32) + + +def _try_keyring_get(account): + """Try to read a secret from the OS keychain. Returns None on failure.""" + try: + import keyring + + return keyring.get_password(_KEYRING_SERVICE, account) + except ImportError: + logger.debug("keyring package not installed") + return None + except Exception as e: + logger.warning("Failed to read from OS keychain: %s", e) + return None + + +def _try_keyring_set(account, value): + """Try to store a secret in the OS keychain. Returns True on success.""" + try: + import keyring + + keyring.set_password(_KEYRING_SERVICE, account, value) + return True + except ImportError: + logger.debug("keyring package not installed") + return False + except Exception as e: + logger.warning("Failed to write to OS keychain: %s", e) + return False + + +def load_or_create_token(config_dir): + """Load auth token from OS keychain, falling back to file. + + Storage priority: + 1. OS Keychain (macOS Keychain / Windows Credential Manager / Linux Secret Service) + 2. File at {config_dir}/auth_token with 0o600 permissions + + Returns the token string. + """ + account = "auth_token" + + # Try keychain first + token = _try_keyring_get(account) + if token: + return token + + # Try file fallback + token_path = Path(config_dir) / "auth_token" + if token_path.exists(): + token = token_path.read_text().strip() + # Migrate to keychain if possible + _try_keyring_set(account, token) + return token + + # Generate new token + token = generate_auth_token() + + # Store in keychain + if _try_keyring_set(account, token): + logger.info("Auth token stored in OS keychain.") + else: + # Fall back to file + logger.warning( + "OS keychain unavailable. Storing auth token in file. " + "Install 'keyring' for secure storage: pip install mempalace[security]", + ) + config_path = Path(config_dir) + config_path.mkdir(parents=True, exist_ok=True) + secure_dir(config_path) + token_path.write_text(token) + secure_file(token_path) + + return token + + +def verify_token(provided, expected): + """Constant-time token comparison to prevent timing attacks.""" + return hmac.compare_digest(provided, expected) + + +# ── Encryption ─────────────────────────────────────────────────────────── + + +def _require_cryptography(): + """Import and return Fernet, raising a clear error if not installed.""" + try: + from cryptography.fernet import Fernet + + return Fernet + except ImportError: + raise ImportError( + "Encryption requires the 'cryptography' package. " + "Install it with: pip install mempalace[security]" + ) + + +def generate_encryption_key(): + """Generate a new Fernet encryption key (bytes).""" + Fernet = _require_cryptography() + return Fernet.generate_key() + + +def load_or_create_key(config_dir): + """Load encryption key from OS keychain, falling back to file. + + Storage priority: + 1. OS Keychain (macOS Keychain / Windows Credential Manager / Linux Secret Service) + 2. File at {config_dir}/palace.key with 0o600 permissions + + Returns a Fernet instance ready for encrypt/decrypt. + """ + Fernet = _require_cryptography() + account = "encryption_key" + + # Try keychain first + key_str = _try_keyring_get(account) + if key_str: + return Fernet(key_str.encode()) + + # Try file fallback + key_path = Path(config_dir) / "palace.key" + if key_path.exists(): + key_bytes = key_path.read_text().strip().encode() + # Migrate to keychain if possible + _try_keyring_set(account, key_bytes.decode()) + return Fernet(key_bytes) + + # Generate new key + key_bytes = generate_encryption_key() + + # Store in keychain + if _try_keyring_set(account, key_bytes.decode()): + logger.info("Encryption key stored in OS keychain.") + else: + # Fall back to file + logger.warning( + "OS keychain unavailable. Storing encryption key in file. " + "Install 'keyring' for secure storage: pip install mempalace[security]", + ) + config_path = Path(config_dir) + config_path.mkdir(parents=True, exist_ok=True) + secure_dir(config_path) + key_path.write_text(key_bytes.decode()) + secure_file(key_path) + + return Fernet(key_bytes) + + +def encrypt(fernet, plaintext): + """Encrypt plaintext string, returning base64-encoded ciphertext string.""" + return fernet.encrypt(plaintext.encode()).decode() + + +def decrypt(fernet, ciphertext): + """Decrypt base64-encoded ciphertext string, returning plaintext string.""" + return fernet.decrypt(ciphertext.encode()).decode() diff --git a/pyproject.toml b/pyproject.toml index 4862873f49..690ed0fd04 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,6 +40,7 @@ mempalace = "mempalace:main" [project.optional-dependencies] dev = ["pytest>=7.0", "ruff>=0.4.0"] spellcheck = ["autocorrect>=2.0"] +security = ["cryptography>=41.0", "keyring>=25.0"] [dependency-groups] dev = ["pytest>=7.0", "ruff>=0.4.0"] diff --git a/tests/test_config.py b/tests/test_config.py index a36b74de22..c27e55d9f2 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -30,3 +30,42 @@ def test_init(): cfg = MempalaceConfig(config_dir=tmpdir) cfg.init() assert os.path.exists(os.path.join(tmpdir, "config.json")) + + +# ── Security config ────────────────────────────────────────────────────── + + +def test_security_config_defaults(): + cfg = MempalaceConfig(config_dir=tempfile.mkdtemp()) + assert cfg.auth_enabled is False + assert cfg.encryption_enabled is False + assert cfg.max_content_size == 1_048_576 + + +def test_security_config_from_file(): + tmpdir = tempfile.mkdtemp() + with open(os.path.join(tmpdir, "config.json"), "w") as f: + json.dump( + { + "security": { + "auth_enabled": True, + "encryption_enabled": True, + "max_content_size": 500_000, + } + }, + f, + ) + cfg = MempalaceConfig(config_dir=tmpdir) + assert cfg.auth_enabled is True + assert cfg.encryption_enabled is True + assert cfg.max_content_size == 500_000 + + +def test_security_env_override(): + tmpdir = tempfile.mkdtemp() + cfg = MempalaceConfig(config_dir=tmpdir) + os.environ["MEMPALACE_AUTH_ENABLED"] = "true" + try: + assert cfg.auth_enabled is True + finally: + del os.environ["MEMPALACE_AUTH_ENABLED"] diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index cf37a27709..ea4088a171 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -8,6 +8,8 @@ import json +import pytest + def _patch_mcp_server(monkeypatch, config, palace_path, kg): """Patch the mcp_server module globals to use test fixtures.""" @@ -336,3 +338,293 @@ def test_diary_read_empty(self, monkeypatch, config, palace_path, kg): r = tool_diary_read(agent_name="Nobody") assert r["entries"] == [] + + +# ── Input Validation ─────────────────────────────────────────────────── + + +class TestLocalhostGuard: + def test_start_http_server_rejects_public_address(self): + from mempalace.mcp_server import start_http_server + + with pytest.raises(ValueError, match="non-localhost"): + start_http_server(host="0.0.0.0", port=8766) + + def test_start_http_server_accepts_localhost(self): + from mempalace.mcp_server import start_http_server + + # Should pass validation but raise NotImplementedError (HTTP not built yet) + with pytest.raises(NotImplementedError): + start_http_server(host="127.0.0.1", port=8766) + + +class TestAuth: + def test_auth_disabled_no_token_needed(self, monkeypatch, config, palace_path, kg): + _patch_mcp_server(monkeypatch, config, palace_path, kg) + from mempalace import mcp_server + from mempalace.mcp_server import handle_request + + monkeypatch.setattr(mcp_server, "_auth_token", None) + _get_collection(palace_path, create=True) + + resp = handle_request( + { + "method": "tools/call", + "id": 200, + "params": {"name": "mempalace_status", "arguments": {}}, + } + ) + assert "result" in resp + + def test_auth_enabled_missing_token(self, monkeypatch, config, palace_path, kg): + _patch_mcp_server(monkeypatch, config, palace_path, kg) + from mempalace import mcp_server + from mempalace.mcp_server import handle_request + + monkeypatch.setattr(mcp_server, "_auth_token", "secret-token-123") + + resp = handle_request( + { + "method": "tools/call", + "id": 201, + "params": {"name": "mempalace_status", "arguments": {}}, + } + ) + assert resp["error"]["code"] == -32001 + + def test_auth_enabled_wrong_token(self, monkeypatch, config, palace_path, kg): + _patch_mcp_server(monkeypatch, config, palace_path, kg) + from mempalace import mcp_server + from mempalace.mcp_server import handle_request + + monkeypatch.setattr(mcp_server, "_auth_token", "secret-token-123") + + resp = handle_request( + { + "method": "tools/call", + "id": 202, + "params": { + "name": "mempalace_status", + "arguments": {}, + "_meta": {"auth_token": "wrong-token"}, + }, + } + ) + assert resp["error"]["code"] == -32001 + + def test_auth_enabled_valid_token(self, monkeypatch, config, palace_path, kg): + _patch_mcp_server(monkeypatch, config, palace_path, kg) + from mempalace import mcp_server + from mempalace.mcp_server import handle_request + + monkeypatch.setattr(mcp_server, "_auth_token", "secret-token-123") + _get_collection(palace_path, create=True) + + resp = handle_request( + { + "method": "tools/call", + "id": 203, + "params": { + "name": "mempalace_status", + "arguments": {}, + "_meta": {"auth_token": "secret-token-123"}, + }, + } + ) + assert "result" in resp + + def test_initialize_reports_auth_required(self, monkeypatch): + from mempalace import mcp_server + from mempalace.mcp_server import handle_request + + monkeypatch.setattr(mcp_server, "_auth_token", "some-token") + + resp = handle_request({"method": "initialize", "id": 204, "params": {}}) + assert resp["result"]["serverInfo"]["authRequired"] is True + + def test_initialize_no_auth_required_when_disabled(self, monkeypatch): + from mempalace import mcp_server + from mempalace.mcp_server import handle_request + + monkeypatch.setattr(mcp_server, "_auth_token", None) + + resp = handle_request({"method": "initialize", "id": 205, "params": {}}) + assert "authRequired" not in resp["result"]["serverInfo"] + + def test_tools_list_requires_auth(self, monkeypatch, config, palace_path, kg): + _patch_mcp_server(monkeypatch, config, palace_path, kg) + from mempalace import mcp_server + from mempalace.mcp_server import handle_request + + monkeypatch.setattr(mcp_server, "_auth_token", "secret-token-123") + + resp = handle_request({"method": "tools/list", "id": 206, "params": {}}) + assert resp["error"]["code"] == -32001 + + +class TestInputValidation: + def test_missing_required_field(self, monkeypatch, config, palace_path, kg): + _patch_mcp_server(monkeypatch, config, palace_path, kg) + from mempalace.mcp_server import handle_request + + resp = handle_request( + { + "method": "tools/call", + "id": 100, + "params": { + "name": "mempalace_search", + "arguments": {}, # missing required "query" + }, + } + ) + assert resp["error"]["code"] == -32602 + assert "query" in resp["error"]["message"] + + def test_extra_unknown_field_rejected(self, monkeypatch, config, palace_path, kg): + _patch_mcp_server(monkeypatch, config, palace_path, kg) + from mempalace.mcp_server import handle_request + + resp = handle_request( + { + "method": "tools/call", + "id": 101, + "params": { + "name": "mempalace_search", + "arguments": {"query": "test", "bogus_field": "nope"}, + }, + } + ) + assert resp["error"]["code"] == -32602 + assert "bogus_field" in resp["error"]["message"] + + def test_content_size_limit_exceeded(self, monkeypatch, config, palace_path, kg): + _patch_mcp_server(monkeypatch, config, palace_path, kg) + # Set a very small limit for testing + config._file_config["security"] = {"max_content_size": 100} + from mempalace.mcp_server import handle_request + + resp = handle_request( + { + "method": "tools/call", + "id": 102, + "params": { + "name": "mempalace_add_drawer", + "arguments": { + "wing": "test", + "room": "test", + "content": "x" * 200, + }, + }, + } + ) + assert resp["error"]["code"] == -32000 + assert "max size" in resp["error"]["message"] + + def test_content_within_size_limit(self, monkeypatch, config, palace_path, kg): + _patch_mcp_server(monkeypatch, config, palace_path, kg) + _get_collection(palace_path, create=True) + from mempalace.mcp_server import handle_request + + resp = handle_request( + { + "method": "tools/call", + "id": 103, + "params": { + "name": "mempalace_add_drawer", + "arguments": { + "wing": "test", + "room": "test", + "content": "Small content that fits.", + }, + }, + } + ) + assert "result" in resp + + +# ── Encryption ───────────────────────────────────────────────────────── + + +class TestEncryption: + def _enable_encryption(self, monkeypatch): + from cryptography.fernet import Fernet + from mempalace import mcp_server + + f = Fernet(Fernet.generate_key()) + monkeypatch.setattr(mcp_server, "_fernet", f) + return f + + def test_add_drawer_stores_encrypted_content(self, monkeypatch, config, palace_path, kg): + _patch_mcp_server(monkeypatch, config, palace_path, kg) + f = self._enable_encryption(monkeypatch) + col = _get_collection(palace_path, create=True) + from mempalace.mcp_server import tool_add_drawer + + result = tool_add_drawer( + wing="enc_test", + room="secrets", + content="Top secret family information.", + ) + assert result["success"] is True + + # Verify encrypted_content is in metadata and plaintext is NOT in documents + stored = col.get(ids=[result["drawer_id"]], include=["metadatas", "documents"]) + meta = stored["metadatas"][0] + doc = stored["documents"][0] + assert "encrypted_content" in meta + assert meta["encrypted_content"] != "Top secret family information." + # Documents field should contain a placeholder, not plaintext + assert "Top secret family information." not in doc + assert doc.startswith("[encrypted:") + + # Verify it decrypts correctly + from mempalace.security import decrypt + + assert decrypt(f, meta["encrypted_content"]) == "Top secret family information." + + def test_search_returns_decrypted_content(self, monkeypatch, config, palace_path, kg): + _patch_mcp_server(monkeypatch, config, palace_path, kg) + self._enable_encryption(monkeypatch) + _get_collection(palace_path, create=True) + from mempalace.mcp_server import tool_add_drawer, tool_search + + tool_add_drawer( + wing="enc_test", + room="secrets", + content="The encryption key is stored in the OS keychain for safety.", + ) + + result = tool_search(query="encryption key keychain") + assert len(result["results"]) > 0 + top = result["results"][0] + assert "encryption key" in top["text"].lower() or "keychain" in top["text"].lower() + + def test_diary_write_read_encrypted_roundtrip(self, monkeypatch, config, palace_path, kg): + _patch_mcp_server(monkeypatch, config, palace_path, kg) + self._enable_encryption(monkeypatch) + _get_collection(palace_path, create=True) + from mempalace.mcp_server import tool_diary_read, tool_diary_write + + w = tool_diary_write( + agent_name="SecureAgent", + entry="Encrypted diary: today we added Fernet encryption.", + topic="security", + ) + assert w["success"] is True + + r = tool_diary_read(agent_name="SecureAgent") + assert r["total"] == 1 + assert "Fernet encryption" in r["entries"][0]["content"] + + def test_unencrypted_drawers_still_readable( + self, monkeypatch, config, palace_path, seeded_collection, kg + ): + """Drawers added before encryption was enabled should still be readable.""" + _patch_mcp_server(monkeypatch, config, palace_path, kg) + self._enable_encryption(monkeypatch) + from mempalace.mcp_server import tool_search + + # seeded_collection has drawers without encrypted_content + result = tool_search(query="JWT authentication") + assert len(result["results"]) > 0 + assert "JWT" in result["results"][0]["text"] diff --git a/tests/test_security.py b/tests/test_security.py new file mode 100644 index 0000000000..198b44b101 --- /dev/null +++ b/tests/test_security.py @@ -0,0 +1,228 @@ +""" +test_security.py — Tests for the security module. +""" + +import os +import stat +import sys +import tempfile + +import pytest + +from mempalace.security import ( + content_hash, + decrypt, + encrypt, + generate_auth_token, + generate_encryption_key, + load_or_create_key, + load_or_create_token, + secure_dir, + secure_file, + validate_bind_address, + verify_token, +) + + +# ── Hashing ────────────────────────────────────────────────────────────── + + +class TestContentHash: + def test_returns_hex_string(self): + result = content_hash("hello world") + assert all(c in "0123456789abcdef" for c in result) + + def test_default_length_is_16(self): + assert len(content_hash("test data")) == 16 + + def test_custom_length(self): + assert len(content_hash("test data", length=8)) == 8 + assert len(content_hash("test data", length=32)) == 32 + + def test_deterministic(self): + assert content_hash("same input") == content_hash("same input") + + def test_different_inputs_differ(self): + assert content_hash("input a") != content_hash("input b") + + def test_uses_sha256(self): + """Verify the output matches a known SHA-256 prefix.""" + import hashlib + + expected = hashlib.sha256(b"test").hexdigest()[:16] + assert content_hash("test") == expected + + +# ── File Permissions ───────────────────────────────────────────────────── + + +@pytest.mark.skipif(sys.platform == "win32", reason="Unix permissions") +class TestFilePermissions: + def test_secure_file_sets_600(self): + with tempfile.NamedTemporaryFile(delete=False) as f: + f.write(b"secret") + path = f.name + try: + secure_file(path) + mode = stat.S_IMODE(os.stat(path).st_mode) + assert mode == 0o600 + finally: + os.unlink(path) + + def test_secure_dir_sets_700(self): + d = tempfile.mkdtemp() + try: + secure_dir(d) + mode = stat.S_IMODE(os.stat(d).st_mode) + assert mode == 0o700 + finally: + os.rmdir(d) + + def test_secure_file_nonexistent_does_not_raise(self): + """Should log a warning but not crash.""" + secure_file("/nonexistent/path/to/file") + + def test_secure_dir_nonexistent_does_not_raise(self): + """Should log a warning but not crash.""" + secure_dir("/nonexistent/path/to/dir") + + +# ── Localhost Validation ───────────────────────────────────────────────── + + +class TestValidateBindAddress: + def test_allows_127_0_0_1(self): + validate_bind_address("127.0.0.1") # should not raise + + def test_allows_ipv6_loopback(self): + validate_bind_address("::1") # should not raise + + def test_allows_localhost(self): + validate_bind_address("localhost") # should not raise + + def test_rejects_0_0_0_0(self): + with pytest.raises(ValueError, match="non-localhost"): + validate_bind_address("0.0.0.0") + + def test_rejects_public_ip(self): + with pytest.raises(ValueError, match="non-localhost"): + validate_bind_address("192.168.1.1") + + def test_rejects_wildcard_ipv6(self): + with pytest.raises(ValueError, match="non-localhost"): + validate_bind_address("::") + + def test_rejects_empty_string(self): + with pytest.raises(ValueError, match="non-localhost"): + validate_bind_address("") + + +# ── Authentication ─────────────────────────────────────────────────────── + + +class TestAuthToken: + def test_generate_auth_token_format(self): + token = generate_auth_token() + assert isinstance(token, str) + assert len(token) == 43 # secrets.token_urlsafe(32) produces 43 chars + + def test_generate_auth_token_unique(self): + assert generate_auth_token() != generate_auth_token() + + def test_verify_token_correct(self): + assert verify_token("my-secret", "my-secret") is True + + def test_verify_token_wrong(self): + assert verify_token("wrong", "my-secret") is False + + def test_verify_token_empty(self): + assert verify_token("", "my-secret") is False + + def test_load_or_create_token_creates(self, monkeypatch): + """Token is created and returned on first call (file fallback).""" + # Force keyring to fail so we test file fallback + monkeypatch.setattr("mempalace.security._try_keyring_get", lambda account: None) + monkeypatch.setattr("mempalace.security._try_keyring_set", lambda account, value: False) + d = tempfile.mkdtemp() + token = load_or_create_token(d) + assert isinstance(token, str) + assert len(token) == 43 + # Token file exists with restrictive permissions + token_path = os.path.join(d, "auth_token") + assert os.path.exists(token_path) + if sys.platform != "win32": + mode = stat.S_IMODE(os.stat(token_path).st_mode) + assert mode == 0o600 + + def test_load_or_create_token_reads_existing(self, monkeypatch): + """Second call returns the same token.""" + monkeypatch.setattr("mempalace.security._try_keyring_get", lambda account: None) + monkeypatch.setattr("mempalace.security._try_keyring_set", lambda account, value: False) + d = tempfile.mkdtemp() + token1 = load_or_create_token(d) + token2 = load_or_create_token(d) + assert token1 == token2 + + +# ── Encryption ─────────────────────────────────────────────────────────── + + +class TestEncryption: + def test_generate_encryption_key(self): + key = generate_encryption_key() + assert isinstance(key, bytes) + assert len(key) == 44 # Fernet key is 44 bytes base64-encoded + + def test_encrypt_decrypt_roundtrip(self): + from cryptography.fernet import Fernet + + f = Fernet(Fernet.generate_key()) + plaintext = "This is sensitive palace data about family relationships." + ciphertext = encrypt(f, plaintext) + assert ciphertext != plaintext + assert decrypt(f, ciphertext) == plaintext + + def test_encrypt_produces_different_ciphertext(self): + """Fernet uses random IV, so two encryptions of the same text differ.""" + from cryptography.fernet import Fernet + + f = Fernet(Fernet.generate_key()) + plaintext = "same input" + ct1 = encrypt(f, plaintext) + ct2 = encrypt(f, plaintext) + assert ct1 != ct2 + # Both decrypt to same plaintext + assert decrypt(f, ct1) == plaintext + assert decrypt(f, ct2) == plaintext + + def test_decrypt_wrong_key_fails(self): + from cryptography.fernet import Fernet, InvalidToken + + f1 = Fernet(Fernet.generate_key()) + f2 = Fernet(Fernet.generate_key()) + ciphertext = encrypt(f1, "secret") + with pytest.raises(InvalidToken): + decrypt(f2, ciphertext) + + def test_load_or_create_key_creates(self, monkeypatch): + """Key is created and returned on first call (file fallback).""" + monkeypatch.setattr("mempalace.security._try_keyring_get", lambda account: None) + monkeypatch.setattr("mempalace.security._try_keyring_set", lambda account, value: False) + d = tempfile.mkdtemp() + fernet = load_or_create_key(d) + # Verify it's a working Fernet instance + ct = encrypt(fernet, "test") + assert decrypt(fernet, ct) == "test" + # Key file exists + assert os.path.exists(os.path.join(d, "palace.key")) + + def test_load_or_create_key_reads_existing(self, monkeypatch): + """Second call returns a Fernet with the same key.""" + monkeypatch.setattr("mempalace.security._try_keyring_get", lambda account: None) + monkeypatch.setattr("mempalace.security._try_keyring_set", lambda account, value: False) + d = tempfile.mkdtemp() + f1 = load_or_create_key(d) + f2 = load_or_create_key(d) + # Both should decrypt each other's ciphertext + ct = encrypt(f1, "cross-test") + assert decrypt(f2, ct) == "cross-test"