Russell Ballestrini
bc11c516b6
New files: - filevault.py: Hash-based file storage with use_pairs option (v1.1.0) - async_filevault.py: Async wrapper using asyncio.to_thread() - domain_vault.py: Triple vault system for web archival (HTML, Media, Linkpeek) - screenshot.py: Async screenshot capture using uri2png - tests/unit/: Comprehensive test suite (85 tests) Sync-to-async conversions: - storage.py: Wrap Path operations in asyncio.to_thread() - domain_vault.py: Wrap exists(), mkdir(), rglob(), os.walk() in asyncio.to_thread() - screenshot.py: Wrap read_bytes(), write_bytes(), unlink() in asyncio.to_thread() All sync filesystem operations now run in thread pool to avoid blocking async loop.
545 lines
18 KiB
Python
545 lines
18 KiB
Python
"""
|
|
FileVault 1.1.0 - Python 2/3 compatible hash-based file storage system
|
|
|
|
A Vault manages a hash directory tree of files on a filesystem.
|
|
|
|
Features:
|
|
- Create hash directory trees of custom depth
|
|
- Spread out files to keep CLI snappy when traversing the tree
|
|
- Scale to hundreds of thousands of files
|
|
- Obfuscate directory paths and filenames
|
|
- Compatible with Python 2.7+ and Python 3.x
|
|
"""
|
|
|
|
from __future__ import unicode_literals, print_function
|
|
|
|
import sys
|
|
import json
|
|
import fcntl
|
|
import contextlib
|
|
import tempfile
|
|
from os import path, makedirs, rename
|
|
from uuid import uuid4
|
|
from hashlib import sha256
|
|
from itertools import permutations
|
|
|
|
# Python 2/3 compatibility
|
|
if sys.version_info[0] == 3:
|
|
string_types = str
|
|
|
|
def ensure_bytes(s):
|
|
if isinstance(s, str):
|
|
return s.encode("utf-8")
|
|
return s
|
|
|
|
else:
|
|
string_types = basestring
|
|
|
|
def ensure_bytes(s):
|
|
if isinstance(s, unicode):
|
|
return s.encode("utf-8")
|
|
return s
|
|
|
|
|
|
__version__ = "1.1.0"
|
|
__author__ = "Russell Ballestrini"
|
|
__email__ = "russell@ballestrini.net"
|
|
__license__ = "Public Domain"
|
|
|
|
HEX = "0123456789abcdef"
|
|
|
|
|
|
class Vault(object):
|
|
"""
|
|
Hash-based file storage system with deterministic directory structure.
|
|
|
|
Creates a directory tree using hex characters for balanced file distribution.
|
|
Includes thread-safe file locking and atomic write operations.
|
|
|
|
Optional in-memory existence cache to avoid repeated filesystem stat() calls.
|
|
|
|
Directory modes:
|
|
- use_pairs=False (default): Single hex chars per level (16 dirs/level)
|
|
Example: depth=3 -> a/b/c/hash
|
|
- use_pairs=True: Hex pairs per level (256 dirs/level, git-style)
|
|
Example: depth=3 -> ab/cd/ef/hash
|
|
"""
|
|
|
|
def __init__(
|
|
self, vaultpath="vault", depth=3, salt="changeme", enable_memory_cache=False,
|
|
use_pairs=False
|
|
):
|
|
"""
|
|
Initialize a new Vault instance.
|
|
|
|
Args:
|
|
vaultpath (str): Base path for vault storage (default: 'vault')
|
|
depth (int): Directory tree depth (default: 3)
|
|
salt (str): Salt for hash generation (default: 'changeme')
|
|
enable_memory_cache (bool): Enable in-memory file existence cache (default: False)
|
|
use_pairs (bool): Use hex pairs for directories instead of single chars (default: False)
|
|
Pairs give 256 subdirs/level vs 16, better for large vaults
|
|
"""
|
|
self.vaultpath = vaultpath
|
|
self.depth = depth
|
|
self.salt = ensure_bytes(salt)
|
|
self.enable_memory_cache = enable_memory_cache
|
|
self.use_pairs = use_pairs
|
|
|
|
# In-memory file existence cache (optional performance optimization)
|
|
self._existence_cache = {} if enable_memory_cache else None
|
|
|
|
self.init_vault()
|
|
|
|
def init_vault(self):
|
|
"""Build the vault base directory if it doesn't exist."""
|
|
if not path.exists(self.vaultpath):
|
|
try:
|
|
makedirs(self.vaultpath)
|
|
except OSError:
|
|
pass
|
|
|
|
def _generate_filename(self, h, ext="", absolute=False):
|
|
"""
|
|
Accept a hash, return a valid file path.
|
|
|
|
Args:
|
|
h (str): Hash string
|
|
ext (str): File extension (optional)
|
|
absolute (bool): Return absolute path if True
|
|
|
|
Returns:
|
|
str: Generated file path
|
|
"""
|
|
if self.use_pairs:
|
|
# Pairs mode: take 2 chars at a time (256 dirs/level, git-style)
|
|
dirs = [h[i:i+2] for i in range(0, self.depth * 2, 2)]
|
|
else:
|
|
# Original mode: single chars (16 dirs/level)
|
|
dirs = [h[i] for i in range(self.depth)]
|
|
if ext and not ext.startswith("."):
|
|
ext = "." + ext
|
|
if absolute:
|
|
return path.join(self.vaultpath, *dirs) + "/" + h + ext
|
|
return path.join(*dirs) + "/" + h + ext
|
|
|
|
def create_filename(self, seed, ext="", absolute=False):
|
|
"""
|
|
Create a valid vault filename seeded with input.
|
|
|
|
Args:
|
|
seed (str): Seed string for deterministic hash
|
|
ext (str): Optional file extension
|
|
absolute (bool): Return absolute path if True
|
|
|
|
Returns:
|
|
str: Generated filename path
|
|
"""
|
|
seed_bytes = ensure_bytes(seed)
|
|
h = sha256(seed_bytes + self.salt).hexdigest()
|
|
return self._generate_filename(h, ext, absolute)
|
|
|
|
def create_random_filename(self, ext="", absolute=False):
|
|
"""
|
|
Create a valid vault filename seeded with random input.
|
|
|
|
Args:
|
|
ext (str): Optional file extension
|
|
absolute (bool): Return absolute path if True
|
|
|
|
Returns:
|
|
str: Generated random filename path
|
|
"""
|
|
random_seed = ensure_bytes(str(uuid4()))
|
|
h = sha256(random_seed + self.salt).hexdigest()
|
|
return self._generate_filename(h, ext, absolute)
|
|
|
|
@contextlib.contextmanager
|
|
def file_lock(self, file_path, mode="r", timeout=30):
|
|
"""
|
|
Context manager for file locking with automatic release.
|
|
|
|
Args:
|
|
file_path (str): Path to file to lock
|
|
mode (str): File open mode ('r', 'w', 'a', etc.)
|
|
timeout (int): Lock timeout in seconds
|
|
|
|
Yields:
|
|
file: Locked file object
|
|
|
|
Raises:
|
|
IOError: If lock cannot be acquired within timeout
|
|
"""
|
|
# Ensure directory exists for write operations
|
|
if "w" in mode or "a" in mode:
|
|
dir_path = path.dirname(file_path)
|
|
if dir_path and not path.exists(dir_path):
|
|
try:
|
|
makedirs(dir_path)
|
|
except OSError:
|
|
pass
|
|
|
|
# Open file and acquire lock
|
|
with open(file_path, mode) as f:
|
|
try:
|
|
# Use exclusive lock for write operations, shared for read
|
|
lock_type = (
|
|
fcntl.LOCK_EX if ("w" in mode or "a" in mode) else fcntl.LOCK_SH
|
|
)
|
|
fcntl.flock(f.fileno(), lock_type | fcntl.LOCK_NB)
|
|
yield f
|
|
except IOError as e:
|
|
if e.errno == 11 or e.errno == 35: # EAGAIN or EWOULDBLOCK
|
|
raise IOError(
|
|
"Could not acquire file lock for: {}".format(file_path)
|
|
)
|
|
raise
|
|
finally:
|
|
try:
|
|
fcntl.flock(f.fileno(), fcntl.LOCK_UN)
|
|
except:
|
|
pass # Ignore unlock errors
|
|
|
|
def atomic_write_json(self, file_path, data, indent=2, aliases=None):
|
|
"""
|
|
Atomically write JSON data to a file with proper locking.
|
|
|
|
Uses a temporary file + rename for atomic operations and fcntl for locking.
|
|
|
|
Args:
|
|
file_path (str): Target file path
|
|
data (dict/list): Data to write as JSON
|
|
indent (int): JSON indentation level
|
|
aliases (list): Optional list of alias keys that will create symlinks to this file
|
|
|
|
Raises:
|
|
IOError: If file operations fail
|
|
ValueError: If data cannot be serialized to JSON
|
|
"""
|
|
# Ensure target directory exists
|
|
dir_path = path.dirname(file_path)
|
|
if dir_path and not path.exists(dir_path):
|
|
try:
|
|
makedirs(dir_path)
|
|
except OSError:
|
|
pass
|
|
|
|
# Create temporary file in same directory for atomic rename
|
|
temp_fd, temp_path = tempfile.mkstemp(
|
|
dir=dir_path or ".",
|
|
prefix=".tmp_" + path.basename(file_path) + "_",
|
|
suffix=".json",
|
|
)
|
|
|
|
try:
|
|
with open(temp_path, "w") as temp_file:
|
|
# Lock the temporary file
|
|
fcntl.flock(temp_file.fileno(), fcntl.LOCK_EX)
|
|
|
|
# Write JSON data
|
|
json.dump(data, temp_file, indent=indent)
|
|
temp_file.flush()
|
|
|
|
# Force write to disk
|
|
import os
|
|
|
|
os.fsync(temp_file.fileno())
|
|
|
|
# Atomic rename - this is the commit point
|
|
rename(temp_path, file_path)
|
|
|
|
# Mark file as existing in cache after successful write
|
|
self.mark_file_exists(file_path, True)
|
|
|
|
# Create aliases (symlinks) if requested
|
|
if aliases:
|
|
for alias_key in aliases:
|
|
self._create_alias(file_path, alias_key)
|
|
|
|
except Exception as e:
|
|
# Clean up temporary file on any error
|
|
try:
|
|
import os
|
|
|
|
os.unlink(temp_path)
|
|
except:
|
|
pass
|
|
raise IOError("Failed to write JSON file {}: {}".format(file_path, str(e)))
|
|
finally:
|
|
# Close temp file descriptor
|
|
try:
|
|
import os
|
|
|
|
os.close(temp_fd)
|
|
except:
|
|
pass
|
|
|
|
def _create_alias(self, target_path, alias_key):
|
|
"""Create a symlink from alias_key to target_path"""
|
|
alias_path = self.create_filename(alias_key, ".json", absolute=True)
|
|
|
|
# Ensure alias directory exists
|
|
alias_dir = path.dirname(alias_path)
|
|
if alias_dir and not path.exists(alias_dir):
|
|
try:
|
|
makedirs(alias_dir)
|
|
except OSError:
|
|
pass
|
|
|
|
# Create or update symlink
|
|
rel_path = None # Initialize to avoid NameError in exception handler
|
|
try:
|
|
# Check if symlink already exists
|
|
if path.exists(alias_path) or path.islink(alias_path):
|
|
# Check if it points to the correct target
|
|
try:
|
|
if path.realpath(alias_path) == path.realpath(target_path):
|
|
# Already points to correct file, nothing to do
|
|
return
|
|
except:
|
|
pass
|
|
# Remove incorrect symlink
|
|
import os
|
|
|
|
os.unlink(alias_path)
|
|
|
|
# Create relative symlink so it works across directory structures
|
|
import os
|
|
|
|
# Calculate relative path from alias to target
|
|
alias_dir = path.dirname(alias_path)
|
|
rel_path = path.relpath(target_path, alias_dir)
|
|
os.symlink(rel_path, alias_path)
|
|
except Exception as e:
|
|
# Log but don't fail the write operation
|
|
import sys
|
|
import traceback
|
|
|
|
print(
|
|
"Warning: Failed to create alias {}: {}".format(alias_key, e),
|
|
file=sys.stderr,
|
|
)
|
|
print("Alias path: {}".format(alias_path), file=sys.stderr)
|
|
print("Target path: {}".format(target_path), file=sys.stderr)
|
|
if rel_path is not None:
|
|
print("Relative path: {}".format(rel_path), file=sys.stderr)
|
|
else:
|
|
print(
|
|
"Relative path: (not calculated due to earlier error)",
|
|
file=sys.stderr,
|
|
)
|
|
traceback.print_exc(file=sys.stderr)
|
|
|
|
def safe_read_json(self, file_path, default=None):
|
|
"""
|
|
Safely read JSON data from a file with proper locking.
|
|
|
|
Args:
|
|
file_path (str): Path to JSON file
|
|
default (any): Default value if file doesn't exist or is invalid
|
|
|
|
Returns:
|
|
dict/list: Parsed JSON data or default value
|
|
|
|
Raises:
|
|
IOError: If file cannot be read
|
|
ValueError: If JSON is malformed and no default provided
|
|
"""
|
|
if not path.exists(file_path):
|
|
if default is not None:
|
|
return default
|
|
raise IOError("File not found: {}".format(file_path))
|
|
|
|
try:
|
|
with self.file_lock(file_path, "r") as f:
|
|
return json.load(f)
|
|
except (IOError, ValueError) as e:
|
|
if default is not None:
|
|
return default
|
|
raise
|
|
|
|
def file_exists(self, file_path):
|
|
"""
|
|
Check if file exists, using memory cache if enabled.
|
|
|
|
Args:
|
|
file_path (str): Path to check
|
|
|
|
Returns:
|
|
bool: True if file exists, False otherwise
|
|
"""
|
|
if self.enable_memory_cache and self._existence_cache is not None:
|
|
# Check memory cache first
|
|
if file_path in self._existence_cache:
|
|
return self._existence_cache[file_path]
|
|
|
|
# Cache miss - check filesystem and cache result
|
|
exists = path.exists(file_path)
|
|
self._existence_cache[file_path] = exists
|
|
return exists
|
|
else:
|
|
# No cache - direct filesystem check
|
|
return path.exists(file_path)
|
|
|
|
def mark_file_exists(self, file_path, exists=True):
|
|
"""
|
|
Mark a file as existing or not existing in the cache.
|
|
|
|
This should be called after file operations to keep cache consistent.
|
|
|
|
Args:
|
|
file_path (str): Path to mark
|
|
exists (bool): Whether file exists (default: True)
|
|
"""
|
|
if self.enable_memory_cache and self._existence_cache is not None:
|
|
self._existence_cache[file_path] = exists
|
|
|
|
def clear_existence_cache(self):
|
|
"""
|
|
Clear the entire existence cache.
|
|
|
|
Useful for testing or when filesystem state may have changed externally.
|
|
"""
|
|
if self.enable_memory_cache and self._existence_cache is not None:
|
|
self._existence_cache.clear()
|
|
|
|
def get_cache_stats(self):
|
|
"""
|
|
Get statistics about the existence cache.
|
|
|
|
Returns:
|
|
dict: Cache statistics or None if cache disabled
|
|
"""
|
|
if self.enable_memory_cache and self._existence_cache is not None:
|
|
return {
|
|
"enabled": True,
|
|
"entries": len(self._existence_cache),
|
|
"hits": sum(1 for exists in self._existence_cache.values() if exists),
|
|
"misses": sum(
|
|
1 for exists in self._existence_cache.values() if not exists
|
|
),
|
|
}
|
|
return {"enabled": False}
|
|
|
|
def remove_file(self, file_path):
|
|
"""
|
|
Remove a file and update the memory cache.
|
|
|
|
Args:
|
|
file_path (str): Path to file to remove
|
|
|
|
Returns:
|
|
bool: True if file was removed, False if it didn't exist
|
|
|
|
Raises:
|
|
OSError: If file removal fails
|
|
"""
|
|
if not path.exists(file_path):
|
|
# Mark as not existing in cache even if file doesn't exist
|
|
self.mark_file_exists(file_path, False)
|
|
return False
|
|
|
|
try:
|
|
import os
|
|
|
|
os.remove(file_path)
|
|
# Mark as not existing in cache after successful removal
|
|
self.mark_file_exists(file_path, False)
|
|
return True
|
|
except OSError as e:
|
|
# Re-raise the error but don't update cache if removal failed
|
|
raise OSError("Failed to remove file {}: {}".format(file_path, str(e)))
|
|
|
|
def remove_file_if_exists(self, file_path):
|
|
"""
|
|
Remove a file if it exists, ignoring errors if file doesn't exist.
|
|
|
|
Args:
|
|
file_path (str): Path to file to remove
|
|
|
|
Returns:
|
|
bool: True if file was removed, False if it didn't exist
|
|
"""
|
|
try:
|
|
return self.remove_file(file_path)
|
|
except OSError:
|
|
# File didn't exist or couldn't be removed, mark as not existing
|
|
self.mark_file_exists(file_path, False)
|
|
return False
|
|
|
|
def purge_cache_for_path_pattern(self, path_pattern):
|
|
"""
|
|
Remove all cache entries matching a path pattern.
|
|
|
|
Useful for GDPR compliance when removing user data.
|
|
|
|
Args:
|
|
path_pattern (str): Pattern to match (simple string contains check)
|
|
"""
|
|
if self.enable_memory_cache and self._existence_cache is not None:
|
|
# Find all cache keys that contain the pattern
|
|
keys_to_remove = [
|
|
key for key in self._existence_cache.keys() if path_pattern in key
|
|
]
|
|
|
|
# Remove matching entries
|
|
for key in keys_to_remove:
|
|
del self._existence_cache[key]
|
|
|
|
return len(keys_to_remove)
|
|
return 0
|
|
|
|
def write_text_file(self, file_path, content, mode="w"):
|
|
"""
|
|
Write text content to a file with proper locking.
|
|
|
|
Args:
|
|
file_path (str): Target file path
|
|
content (str): Text content to write
|
|
mode (str): File open mode ('w', 'a', etc.)
|
|
|
|
Raises:
|
|
IOError: If file operations fail
|
|
"""
|
|
with self.file_lock(file_path, mode) as f:
|
|
f.write(content)
|
|
f.flush()
|
|
import os
|
|
|
|
os.fsync(f.fileno())
|
|
|
|
# Mark file as existing in cache after successful write
|
|
self.mark_file_exists(file_path, True)
|
|
|
|
def read_text_file(self, file_path, default=None):
|
|
"""
|
|
Read text content from a file with proper locking.
|
|
|
|
Args:
|
|
file_path (str): Path to text file
|
|
default (str): Default value if file doesn't exist
|
|
|
|
Returns:
|
|
str: File content or default value
|
|
|
|
Raises:
|
|
IOError: If file cannot be read and no default provided
|
|
"""
|
|
if not path.exists(file_path):
|
|
if default is not None:
|
|
return default
|
|
raise IOError("File not found: {}".format(file_path))
|
|
|
|
try:
|
|
with self.file_lock(file_path, "r") as f:
|
|
return f.read()
|
|
except IOError:
|
|
if default is not None:
|
|
return default
|
|
raise
|
|
|
|
|
|
# For backwards compatibility
|
|
def create_vault(*args, **kwargs):
|
|
"""Factory function for creating Vault instances (backwards compatibility)."""
|
|
return Vault(*args, **kwargs)
|