"""
OMEGA Crypto -- Optional encryption at rest for memory stores.
When enabled, all graph JSON and JSONL entries are encrypted with Fernet
(AES-128-CBC + HMAC-SHA256). The encryption key is derived from a
machine-specific secret stored at ~/.omega/.key.
Enabled by default since v1.1.
Disable: OMEGA_ENCRYPT=0
The key file is created automatically on first use. Guard it with
filesystem permissions (0600). Losing it means losing access to
encrypted memories.
"""
import base64
import logging
import os
import secrets
from pathlib import Path
logger = logging.getLogger("omega.crypto")
_fernet_instance = None
_checked = False
def _omega_home() -> Path:
"""Resolve OMEGA_HOME lazily so tests can override via env var."""
return Path(os.environ.get("OMEGA_HOME", str(Path.home() / ".omega")))
def _key_path() -> Path:
"""Resolve key file path lazily."""
return _omega_home() / ".key"
def is_enabled() -> bool:
"""Check if encryption at rest is enabled (on by default since v1.1).
Set OMEGA_ENCRYPT=0 to explicitly disable.
"""
val = os.environ.get("OMEGA_ENCRYPT", "").strip().lower()
if val in ("0", "false", "no"):
return False
return True # Default: enabled
def reset_crypto_state() -> None:
"""Reset module state for test isolation."""
global _fernet_instance, _checked
_fernet_instance = None
_checked = False
def _get_or_create_key() -> bytes:
"""Get the Fernet key, creating one if it doesn't exist."""
kp = _key_path()
if kp.exists():
raw = kp.read_bytes().strip()
# If it's a 32-byte raw secret, derive Fernet key from it
if len(raw) == 32:
return base64.urlsafe_b64encode(raw)
# Already a base64-encoded Fernet key
return raw
# Generate a new 32-byte secret and derive Fernet key
omega_home = _omega_home()
omega_home.mkdir(parents=True, exist_ok=True, mode=0o700)
raw_secret = secrets.token_bytes(32)
encoded_key = base64.urlsafe_b64encode(raw_secret)
# Atomic creation with restricted permissions (no TOCTOU window)
fd = os.open(str(kp), os.O_CREAT | os.O_WRONLY | os.O_EXCL, 0o600)
try:
os.write(fd, encoded_key)
finally:
os.close(fd)
logger.info(f"Created encryption key at {kp}")
return encoded_key
def _get_fernet():
"""Lazy-load Fernet instance."""
global _fernet_instance, _checked
if _fernet_instance is not None:
return _fernet_instance
if _checked:
return None
_checked = True
try:
from cryptography.fernet import Fernet
except ImportError:
logger.warning("OMEGA_ENCRYPT=1 but 'cryptography' package not installed. Run: pip install cryptography")
return None
try:
key = _get_or_create_key()
_fernet_instance = Fernet(key)
return _fernet_instance
except Exception as e:
logger.error(f"Failed to initialize encryption: {e}", exc_info=True)
return None
def encrypt(plaintext: str) -> str:
"""Encrypt a string. Returns base64-encoded ciphertext.
If encryption is disabled or unavailable, returns plaintext unchanged.
"""
if not is_enabled():
return plaintext
f = _get_fernet()
if f is None:
return plaintext
token = f.encrypt(plaintext.encode("utf-8"))
return "ENC:" + token.decode("ascii")
def decrypt(data: str) -> str:
"""Decrypt a string. Handles both encrypted and plaintext inputs.
If the data doesn't start with 'ENC:', it's returned as-is (plaintext).
This allows transparent migration: old plaintext files are readable,
new writes are encrypted.
Raises ValueError if decryption fails (bad key, corrupted data, or
cryptography package missing) so callers get a clear signal instead
of receiving raw ciphertext.
"""
if not data.startswith("ENC:"):
return data
f = _get_fernet()
if f is None:
raise ValueError(
"Cannot decrypt: cryptography package not available. "
"Install with: pip install cryptography"
)
try:
token = data[4:].encode("ascii")
return f.decrypt(token).decode("utf-8")
except Exception as e:
raise ValueError(f"Decryption failed: {e}") from e
def encrypt_line(line: str) -> str:
"""Encrypt a single JSONL line."""
if not is_enabled():
return line
return encrypt(line)
def decrypt_line(line: str) -> str:
"""Decrypt a single JSONL line. Handles plaintext transparently."""
return decrypt(line.strip())
def secure_connect(db_path, **kwargs):
"""Create a SQLite connection with secure file permissions (0o600).
Pre-creates the DB file with restricted permissions before connecting,
and fixes existing files that have overly permissive permissions.
"""
import sqlite3
import stat
db_path_str = str(db_path)
path_obj = Path(db_path_str)
if not path_obj.exists():
# Pre-create with restricted permissions (no TOCTOU window)
fd = os.open(db_path_str, os.O_CREAT | os.O_WRONLY, 0o600)
os.close(fd)
else:
# Fix existing files with overly permissive permissions
current_mode = path_obj.stat().st_mode
if current_mode & (stat.S_IRWXG | stat.S_IRWXO):
os.chmod(db_path_str, 0o600)
return sqlite3.connect(db_path_str, **kwargs)
