engineering/pig.py
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add filevault system with async wrapper and convert sync ops to async

Browse source 
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.
This commit is contained in:
Russell Ballestrini 2025-12-22 20:18:20 -05:00
parent 6dab26d6fa
commit bc11c516b6
13 changed files with 3861 additions and 144 deletions
Download patch file Download diff file Expand all files Collapse all files

3
.gitignore vendored
View file

@ -11,6 +11,9 @@ venv/
vault/
test_vault/
# Vendored dependencies
vendor/
# IDE
.idea/
.vscode/

25
Makefile
View file

@ -1,4 +1,4 @@
.PHONY: venv install test crawl serp clean
.PHONY: venv install test crawl serp clean vendor-uri2png vendor-install
VENV := .venv
PYTHON := $(VENV)/bin/python
@ -24,9 +24,30 @@ clean:
rm -rf $(VENV) __pycache__ *.pyc
rm -rf test_vault test_neopig.db
# Vendor dependencies
vendor-uri2png:
@echo "Fetching uri2png..."
rm -rf vendor/uri2png
mkdir -p vendor
git clone --depth 1 https://github.com/russellballestrini/uri2png vendor/uri2png
rm -rf vendor/uri2png/.git
@echo "uri2png vendored at vendor/uri2png"
vendor-install: vendor-uri2png venv
$(PIP) install -e vendor/uri2png/python
$(PYTHON) -m playwright install
@echo "uri2png Python package installed"
# Combined server (SERP + screenshot)
server: vendor-install
$(PYTHON) serp.py --host 0.0.0.0 --port 8000
# Examples:
# make install - create venv and install deps
# make test - run functional test
# make serp - start SERP server
# make serp - start basic SERP server
# make server - start combined server (SERP + screenshot)
# make crawl ARGS="https://example.com rick morty --mode images"
# make clean - remove venv and test artifacts
# make vendor-uri2png - fetch uri2png into vendor/
# make vendor-install - install uri2png Python package with screenshot support

120
README.md Normal file
View file

@ -0,0 +1,120 @@
# neopig
*Neo Python Image Grabber*
Based on [pig.py](http://russell.ballestrini.net/python-image-grabber-pig-py/) by Russell Ballestrini
neopig is a full-domain async media crawler with:
- **MD5 deduplication** - same content stored once, multiple sources tracked
- **SQLite metadata index** - searchable page context, alt text, titles
- **Content-addressed vault** - files stored by hash, not filename
- **Web SERP interface** - search and browse crawled media
- **Multi-target crawling** - crawl multiple domains concurrently
- **Resume support** - restart crawls without re-downloading
- **Page screenshots** - optional [uri2png](https://github.com/russellballestrini/uri2png) integration
## Setup
```bash
python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt
```
## Usage
### CLI Crawler
```bash
# Crawl single target for images
python neopig.py https://example.com --mode images
# Crawl multiple targets concurrently
python neopig.py https://site1.com https://site2.com https://site3.com --mode images
# Crawl with keywords for tagging
python neopig.py https://example.com -k "landscape" "nature" --mode images
# Crawl all media (images + videos + audio)
python neopig.py https://example.com --mode media
# Limit crawl depth and pages
python neopig.py https://example.com --depth 5 --max-pages 500 --mode images
# Index URLs without downloading
python neopig.py https://example.com --no-download --mode images
# Enable page screenshots (requires uri2png)
python neopig.py https://example.com --mode images --screenshot
# Screenshots with custom viewport
python neopig.py https://example.com --mode images --screenshot --screenshot-width 1920 --screenshot-height 1080
```
### Options
```
positional arguments:
targets Target URI(s) to crawl
optional arguments:
-h, --help show this help message and exit
-k, --keywords Keywords to tag media with
-m, --mode Crawl mode: text, images, videos, media, all
-d, --depth Crawl depth (-1 = unlimited, default: -1)
-p, --max-pages Maximum pages to crawl (-1 = unlimited)
--db Database path (default: neopig.db)
--vault Vault storage path (default: vault)
--no-download Don't download media, just index URLs
-v, --verbose Verbose output
screenshot options (all off by default, requires uri2png):
--screenshot Enable page screenshots
--screenshot-width Viewport width in pixels (default: 1280)
--screenshot-height Viewport height in pixels (default: 1024)
--screenshot-delay Delay after page load in ms (default: 1000)
```
### Web SERP (Search Engine Results Page)
```bash
# Start the web interface
python serp.py --host 0.0.0.0 --port 8000
# With custom db/vault paths
python serp.py --port 8000 --db mydata.db --vault ./myfiles
```
Then visit:
- `http://localhost:8000/` - Search interface
- `http://localhost:8000/live` - Live feed of recent media
- `http://localhost:8000/crawl` - Start new crawls via web UI
## Architecture
```
neopig.py CLI crawler entry point
serp.py FastAPI web interface (search, live feed, crawl UI)
async_web_fetcher.py Async HTTP client with depth crawling
database.py SQLite schema and queries
storage.py Content-addressed vault (MD5 hash storage)
```
## Database Schema
- **media** - Deduplicated content (md5_hash, type, mime, analysis status)
- **media_sources** - All contexts where media was found (page URL, alt text, title)
- **crawl_jobs** - Crawl history and statistics
## Output
Media is stored in a content-addressed vault:
```
vault/
ab/
abcd1234...5678.jpg
cd/
cdef5678...1234.png
```
Files are named by their MD5 hash, organized in 2-character prefix directories.

221
async_filevault.py Normal file
View file

@ -0,0 +1,221 @@
"""
AsyncFileVault 1.1.0 - Async wrapper for FileVault
Provides async versions of FileVault I/O operations using asyncio.to_thread().
CPU-bound operations (hash generation, path creation) remain synchronous.
Usage:
vault = AsyncVault(vaultpath="vault", depth=9, use_pairs=True)
# Sync operations (fast, no I/O)
filename = vault.create_filename("seed", ".json", absolute=True)
# Async operations (file I/O)
await vault.atomic_write_json(filename, {"key": "value"})
data = await vault.safe_read_json(filename)
"""
import asyncio
from functools import partial
from typing import Any, Optional, List
from filevault import Vault, ensure_bytes, create_vault
__version__ = "1.1.0"
__author__ = "Russell Ballestrini"
__license__ = "Public Domain"
class AsyncVault(Vault):
"""
Async wrapper for Vault with non-blocking file I/O.
Inherits all sync methods from Vault but provides async versions
of file operations that use asyncio.to_thread() to avoid blocking.
Sync methods (instant, no await needed):
- create_filename()
- create_random_filename()
- _generate_filename()
- mark_file_exists()
- clear_existence_cache()
- get_cache_stats()
- purge_cache_for_path_pattern()
Async methods (await required):
- atomic_write_json()
- safe_read_json()
- write_text_file()
- read_text_file()
- file_exists()
- remove_file()
- remove_file_if_exists()
"""
def __init__(
self,
vaultpath: str = "vault",
depth: int = 3,
salt: str = "changeme",
enable_memory_cache: bool = False,
use_pairs: bool = False
):
"""
Initialize AsyncVault with same parameters as Vault.
Args:
vaultpath: Base path for vault storage
depth: Directory tree depth
salt: Salt for hash generation
enable_memory_cache: Enable in-memory file existence cache
use_pairs: Use hex pairs for directories (256 dirs/level vs 16)
"""
super().__init__(
vaultpath=vaultpath,
depth=depth,
salt=salt,
enable_memory_cache=enable_memory_cache,
use_pairs=use_pairs
)
async def atomic_write_json(
self,
file_path: str,
data: Any,
indent: int = 2,
aliases: Optional[List[str]] = None
) -> None:
"""
Async version of atomic JSON write.
Args:
file_path: Target file path
data: Data to write as JSON
indent: JSON indentation level
aliases: Optional list of alias keys for symlinks
"""
await asyncio.to_thread(
super().atomic_write_json,
file_path,
data,
indent,
aliases
)
async def safe_read_json(
self,
file_path: str,
default: Any = None
) -> Any:
"""
Async version of safe JSON read.
Args:
file_path: Path to JSON file
default: Default value if file doesn't exist or is invalid
Returns:
Parsed JSON data or default value
"""
return await asyncio.to_thread(
super().safe_read_json,
file_path,
default
)
async def write_text_file(
self,
file_path: str,
content: str,
mode: str = "w"
) -> None:
"""
Async version of text file write.
Args:
file_path: Target file path
content: Text content to write
mode: File open mode ('w', 'a', etc.)
"""
await asyncio.to_thread(
super().write_text_file,
file_path,
content,
mode
)
async def read_text_file(
self,
file_path: str,
default: Optional[str] = None
) -> str:
"""
Async version of text file read.
Args:
file_path: Path to text file
default: Default value if file doesn't exist
Returns:
File content or default value
"""
return await asyncio.to_thread(
super().read_text_file,
file_path,
default
)
async def file_exists(self, file_path: str) -> bool:
"""
Async version of file existence check.
Note: If memory cache is enabled and the path is cached,
this returns immediately from cache without blocking.
Args:
file_path: Path to check
Returns:
True if file exists, False otherwise
"""
# Check cache first (no I/O needed)
if self.enable_memory_cache and self._existence_cache is not None:
if file_path in self._existence_cache:
return self._existence_cache[file_path]
# Cache miss - need to hit filesystem
return await asyncio.to_thread(super().file_exists, file_path)
async def remove_file(self, file_path: str) -> bool:
"""
Async version of file removal.
Args:
file_path: Path to file to remove
Returns:
True if file was removed, False if it didn't exist
"""
# Call parent's sync remove_file directly via Vault class
return await asyncio.to_thread(Vault.remove_file, self, file_path)
async def remove_file_if_exists(self, file_path: str) -> bool:
"""
Async version of remove_file_if_exists.
Args:
file_path: Path to file to remove
Returns:
True if file was removed, False if it didn't exist
"""
try:
return await asyncio.to_thread(Vault.remove_file, self, file_path)
except OSError:
self.mark_file_exists(file_path, False)
return False
def create_async_vault(*args, **kwargs) -> AsyncVault:
"""Factory function for creating AsyncVault instances."""
return AsyncVault(*args, **kwargs)

17
async_web_fetcher.py
View file

@ -1356,7 +1356,8 @@ class AsyncWebFetcher:
should_continue_callback = None, # Deprecated - kept for compatibility
user_query: Optional[str] = None, # Deprecated - kept for compatibility
mode: CrawlMode = CrawlMode.TEXT,
media_callback = None # Callback for discovered media: async fn(media_item: Dict) -> None
media_callback = None, # Callback for discovered media: async fn(media_item: Dict) -> None
page_callback = None, # Callback for page HTML: async fn(url: str, html: str) -> None
) -> List[Dict[str, str]]:
"""
Intelligent keyword-driven crawl strategy with domain prioritization.
@ -1465,6 +1466,13 @@ class AsyncWebFetcher:
"links_from_page": [l['url'] if isinstance(l, dict) else l for l in links] # Store for CSV
}
# Call page callback to archive raw HTML
if page_callback and html:
try:
await page_callback(start_url, html)
except Exception as e:
logger.warning(f"Page callback failed for {start_url}: {e}")
# Extract media if mode requires it
if mode in (CrawlMode.IMAGES, CrawlMode.VIDEOS, CrawlMode.MEDIA, CrawlMode.ALL):
media_items = extract_media_from_html(html, start_url, mode)
@ -1680,6 +1688,13 @@ class AsyncWebFetcher:
"links_from_page": [l['url'] if isinstance(l, dict) else l for l in child_links]
}
# Call page callback to archive raw HTML
if page_callback and html:
try:
await page_callback(link_url, html)
except Exception as e:
logger.warning(f"Page callback failed for {link_url}: {e}")
# Extract media if mode requires it
if mode in (CrawlMode.IMAGES, CrawlMode.VIDEOS, CrawlMode.MEDIA, CrawlMode.ALL):
media_items = extract_media_from_html(html, link_url, mode)

792
domain_vault.py Normal file
View file

@ -0,0 +1,792 @@
#!/usr/bin/env python3
"""
domain_vault.py - Triple filevault system for web archival
Three separate 9-layer deep filevaults, each with git repos per domain.
Directory names are salted hashes (NEOPIG_VAULT_SALT env var) for privacy.
The same domain + salt produces identical hashes across all three vaults.
1. HTML Vault (html_vault/):
html_vault/{9-layers}/{salted_hash}/
.git/ # Git repository (SSH cloneable)
{url_path}/index.html # Rewritten HTML with neopig media paths
{url_path}/index.html.og # Original HTML with original URIs
crawl_log.json # Crawl history
2. Media Vault (media_vault/):
media_vault/{9-layers}/{salted_hash}/
.git/ # Git repository with LFS (SSH cloneable)
{url_path}/image.png # Media files mirroring original paths
index.json # Media index
3. Linkpeek Vault (linkpeek_vault/):
linkpeek_vault/{9-layers}/{salted_hash}/
.git/ # Git repository with LFS (SSH cloneable)
{url_path}/index.png # Page screenshots mirroring URL paths
index.json # Screenshot index
Environment Variables:
NEOPIG_VAULT_SALT: Secret salt for domain hashing (required for privacy)
Features:
- 9-layer deep hash paths for filesystem distribution
- Salted hash directory names (domain name not exposed)
- Git versioning per domain (all vaults)
- Git LFS for media and linkpeek vaults (large files)
- SSH cloneable repos
- Idempotent: only commits when content changes
- .html.og preserves original, .html has rewritten media URIs
- Tree structure mirrors original domain URL paths
"""
import asyncio
import hashlib
import json
import logging
import os
import re
from datetime import datetime, timezone
from pathlib import Path
from typing import Dict, Any, Optional, List, Tuple
from urllib.parse import urlparse, urljoin
import aiofiles
logger = logging.getLogger(__name__)
def get_vault_salt() -> str:
"""Get the vault salt from environment variable."""
return os.environ.get('NEOPIG_VAULT_SALT', '')
def domain_hash(domain: str, salted: bool = False) -> str:
"""
Generate MD5 hash of domain.
Args:
domain: The domain name
salted: If True, include NEOPIG_VAULT_SALT for directory naming
"""
data = domain.lower()
if salted:
salt = get_vault_salt()
data = f"{salt}:{data}"
return hashlib.md5(data.encode()).hexdigest()
def content_hash(content: bytes) -> str:
"""Generate MD5 hash of content for deduplication."""
return hashlib.md5(content).hexdigest()
def get_filevault_path(base_path: Path, domain: str) -> Path:
"""
Generate 9-layer deep path for filevault.
Uses salted hash for both path layers AND final directory name.
The same domain + salt produces the same hash across all three vaults.
Example: example.com -> base/ab/cd/ef/12/34/56/78/9a/bc/abcdef1234567890abcdef1234567890/
"""
h = domain_hash(domain, salted=True)
# 9 layers of 2-char directories
layers = [h[i:i+2] for i in range(0, 18, 2)] # 9 pairs (18 chars)
path = base_path
for layer in layers:
path = path / layer
# Final directory is the full salted hash, not the domain name
return path / h
def url_to_filepath(url: str) -> str:
"""
Convert URL to filesystem path that mirrors the original structure.
Examples:
https://example.com/ -> index.html
https://example.com/about -> about/index.html
https://example.com/blog/post.html -> blog/post.html
https://example.com/images/logo.png -> images/logo.png
"""
parsed = urlparse(url)
path = parsed.path.strip('/')
if not path:
return 'index.html'
# Check if path has an extension
if '.' in path.split('/')[-1]:
return path
else:
# Directory-style URL, add index.html
return f"{path}/index.html"
def extract_media_urls(html: str, base_url: str) -> List[Tuple[str, str]]:
"""
Extract media URLs from HTML.
Returns list of (absolute_url, original_reference) tuples.
"""
media_patterns = [
# Images
r'<img[^>]+src=["\']([^"\']+)["\']',
r'<source[^>]+src=["\']([^"\']+)["\']',
# CSS backgrounds
r'url\(["\']?([^"\')\s]+)["\']?\)',
# Video/Audio
r'<video[^>]+src=["\']([^"\']+)["\']',
r'<audio[^>]+src=["\']([^"\']+)["\']',
# Links to media
r'<link[^>]+href=["\']([^"\']+\.(?:css|ico|png|jpg|jpeg|gif|svg|woff2?|ttf|eot))["\']',
]
found = []
for pattern in media_patterns:
for match in re.finditer(pattern, html, re.IGNORECASE):
original = match.group(1)
# Skip data URIs and anchors
if original.startswith('data:') or original.startswith('#'):
continue
absolute = urljoin(base_url, original)
found.append((absolute, original))
return found
class GitRepo:
"""Async git operations helper."""
def __init__(self, path: Path, use_lfs: bool = False):
self.path = path
self.use_lfs = use_lfs
async def init(self) -> bool:
"""Initialize git repo. Returns True if newly created."""
git_path = self.path / '.git'
if await asyncio.to_thread(git_path.exists):
return False
await asyncio.to_thread(self.path.mkdir, parents=True, exist_ok=True)
await self._run('git', 'init')
await self._run('git', 'config', 'user.email', 'neopig@localhost')
await self._run('git', 'config', 'user.name', 'neopig')
if self.use_lfs:
await self._init_lfs()
logger.info(f"Initialized git repo at {self.path}")
return True
async def _init_lfs(self):
"""Initialize git LFS for media files."""
await self._run('git', 'lfs', 'install', '--local')
# Track common media extensions
extensions = [
'*.png', '*.jpg', '*.jpeg', '*.gif', '*.webp', '*.svg',
'*.mp4', '*.webm', '*.mov', '*.avi',
'*.mp3', '*.wav', '*.ogg', '*.flac',
'*.pdf', '*.zip', '*.tar', '*.gz',
'*.woff', '*.woff2', '*.ttf', '*.eot',
]
for ext in extensions:
await self._run('git', 'lfs', 'track', ext)
await self._run('git', 'add', '.gitattributes')
logger.info(f"Initialized git LFS at {self.path}")
async def _run(self, *args) -> Tuple[bytes, bytes]:
"""Run git command."""
proc = await asyncio.create_subprocess_exec(
*args,
cwd=str(self.path),
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
)
return await proc.communicate()
async def has_changes(self) -> bool:
"""Check if there are uncommitted changes."""
stdout, _ = await self._run('git', 'status', '--porcelain')
return len(stdout.strip()) > 0
async def add_all(self):
"""Stage all changes."""
await self._run('git', 'add', '-A')
async def commit(self, message: str) -> Optional[str]:
"""Commit if changes exist. Returns commit hash or None."""
await self.add_all()
if not await self.has_changes():
return None
stdout, stderr = await self._run('git', 'commit', '-m', message)
stdout, _ = await self._run('git', 'rev-parse', 'HEAD')
commit_hash = stdout.decode().strip()
logger.info(f"Committed {commit_hash[:8]}: {message}")
return commit_hash
async def log(self, limit: int = 10) -> List[Dict[str, str]]:
"""Get recent commit history."""
stdout, _ = await self._run('git', 'log', f'-{limit}', '--format=%H|%s|%ci')
commits = []
for line in stdout.decode().strip().split('\n'):
if '|' in line:
parts = line.split('|', 2)
if len(parts) == 3:
commits.append({
'hash': parts[0],
'message': parts[1],
'date': parts[2],
})
return commits
class DomainHtmlVault:
"""
Git-tracked HTML vault for a single domain.
Stores pages in a tree structure mirroring the original domain:
- {path}/index.html.og = Original HTML
- {path}/index.html = Rewritten with neopig media paths
"""
def __init__(self, vault_base: str, domain: str, media_base_url: str = '/media'):
self.domain = domain.lower()
self.base_path = Path(vault_base)
self.path = get_filevault_path(self.base_path, self.domain)
self.crawl_log_path = self.path / 'crawl_log.json'
self.media_base_url = media_base_url.rstrip('/')
self._git = GitRepo(self.path, use_lfs=False)
self._initialized = False
async def init(self) -> bool:
"""Initialize vault. Returns True if newly created."""
is_new = await self._git.init()
if not await asyncio.to_thread(self.crawl_log_path.exists):
await self._save_crawl_log({'crawls': []})
is_new = True
self._initialized = True
return is_new
async def _load_crawl_log(self) -> Dict[str, Any]:
if not await asyncio.to_thread(self.crawl_log_path.exists):
return {'crawls': []}
async with aiofiles.open(self.crawl_log_path, 'r') as f:
return json.loads(await f.read())
async def _save_crawl_log(self, log: Dict[str, Any]):
async with aiofiles.open(self.crawl_log_path, 'w') as f:
await f.write(json.dumps(log, indent=2))
def get_media_vault_url(self, media_url: str) -> str:
"""
Convert original media URL to neopig media vault path.
Example: https://example.com/images/logo.png
-> /media/{9-layers}/{salted_hash}/images/logo.png
"""
parsed = urlparse(media_url)
media_domain = parsed.netloc.lower()
media_path = parsed.path.strip('/')
# Get salted hash for media domain (same across all vaults)
h = domain_hash(media_domain, salted=True)
layers = '/'.join([h[i:i+2] for i in range(0, 18, 2)])
return f"{self.media_base_url}/{layers}/{h}/{media_path}"
def rewrite_media_urls(self, html: str, base_url: str, media_mappings: Dict[str, str]) -> str:
"""
Rewrite media URLs in HTML to point to neopig media vault.
media_mappings: {original_url: neopig_vault_path}
"""
result = html
for original, neopig_path in media_mappings.items():
# Escape for regex
escaped = re.escape(original)
result = re.sub(escaped, neopig_path, result)
return result
async def archive_page(
self,
url: str,
html: str,
media_mappings: Dict[str, str] = None,
) -> Tuple[bool, str]:
"""
Archive a page's HTML.
Saves both original (.html.og) and rewritten (.html) versions.
Returns (is_changed, content_hash).
"""
if not self._initialized:
await self.init()
filepath = url_to_filepath(url)
og_path = self.path / f"{filepath}.og"
html_path = self.path / filepath
# Create directories
await asyncio.to_thread(og_path.parent.mkdir, parents=True, exist_ok=True)
# Check if content changed
chash = content_hash(html.encode('utf-8'))
is_changed = True
if await asyncio.to_thread(og_path.exists):
async with aiofiles.open(og_path, 'r', encoding='utf-8') as f:
existing = await f.read()
if content_hash(existing.encode('utf-8')) == chash:
is_changed = False
# Save original HTML
async with aiofiles.open(og_path, 'w', encoding='utf-8') as f:
await f.write(html)
# Save rewritten HTML
rewritten = html
if media_mappings:
rewritten = self.rewrite_media_urls(html, url, media_mappings)
async with aiofiles.open(html_path, 'w', encoding='utf-8') as f:
await f.write(rewritten)
return is_changed, chash
async def get_page(self, url: str, original: bool = False) -> Optional[str]:
"""Get archived page HTML. Set original=True for .html.og version."""
filepath = url_to_filepath(url)
if original:
filepath = f"{filepath}.og"
full_path = self.path / filepath
if not await asyncio.to_thread(full_path.exists):
return None
async with aiofiles.open(full_path, 'r', encoding='utf-8') as f:
return await f.read()
async def finish_crawl(self, stats: Dict[str, Any]) -> Optional[str]:
"""Finish crawl, log it, and commit if changes."""
log = await self._load_crawl_log()
log['crawls'].append({
'timestamp': datetime.now(timezone.utc).isoformat(),
'stats': stats,
})
await self._save_crawl_log(log)
pages_changed = stats.get('pages_changed', 0)
if pages_changed > 0:
return await self._git.commit(f"Crawl: {pages_changed} pages changed")
return None
async def get_stats(self) -> Dict[str, Any]:
"""Get vault statistics."""
log = await self._load_crawl_log()
commits = await self._git.log(5)
# Count HTML files (run sync rglob in thread pool)
def _count_pages():
return sum(1 for _ in self.path.rglob('*.html') if not str(_).endswith('.og'))
page_count = await asyncio.to_thread(_count_pages)
return {
'domain': self.domain,
'path': str(self.path),
'pages_count': page_count,
'crawls_count': len(log.get('crawls', [])),
'recent_commits': commits,
'last_crawl': log['crawls'][-1] if log.get('crawls') else None,
}
class DomainMediaVault:
"""
Git LFS-tracked media vault for a single domain.
Stores media files in a tree structure mirroring the original domain.
"""
def __init__(self, vault_base: str, domain: str):
self.domain = domain.lower()
self.base_path = Path(vault_base)
self.path = get_filevault_path(self.base_path, self.domain)
self.index_path = self.path / 'index.json'
self._git = GitRepo(self.path, use_lfs=True)
self._initialized = False
async def init(self) -> bool:
"""Initialize vault. Returns True if newly created."""
is_new = await self._git.init()
if not await asyncio.to_thread(self.index_path.exists):
await self._save_index({'media': {}, 'created': datetime.now(timezone.utc).isoformat()})
is_new = True
self._initialized = True
return is_new
async def _load_index(self) -> Dict[str, Any]:
if not await asyncio.to_thread(self.index_path.exists):
return {'media': {}}
async with aiofiles.open(self.index_path, 'r') as f:
return json.loads(await f.read())
async def _save_index(self, index: Dict[str, Any]):
async with aiofiles.open(self.index_path, 'w') as f:
await f.write(json.dumps(index, indent=2))
async def archive_media(
self,
url: str,
content: bytes,
page_url: str = '',
) -> Tuple[bool, str, Path]:
"""
Archive media content.
Returns (is_new, content_hash, file_path).
"""
if not self._initialized:
await self.init()
# Use URL path as file path
parsed = urlparse(url)
filepath = parsed.path.strip('/')
if not filepath:
# Fallback to content hash
filepath = f"{content_hash(content)}.bin"
file_path = self.path / filepath
await asyncio.to_thread(file_path.parent.mkdir, parents=True, exist_ok=True)
# Check if new or changed
chash = content_hash(content)
exists = await asyncio.to_thread(file_path.exists)
is_new = not exists
if exists:
async with aiofiles.open(file_path, 'rb') as f:
existing_hash = content_hash(await f.read())
is_new = existing_hash != chash
# Save media
async with aiofiles.open(file_path, 'wb') as f:
await f.write(content)
# Update index
index = await self._load_index()
index['media'][filepath] = {
'url': url,
'hash': chash,
'size': len(content),
'page_url': page_url,
'updated': datetime.now(timezone.utc).isoformat(),
}
await self._save_index(index)
return is_new, chash, file_path
async def get_media(self, url: str) -> Optional[Tuple[bytes, Dict[str, Any]]]:
"""Get archived media content and metadata."""
parsed = urlparse(url)
filepath = parsed.path.strip('/')
file_path = self.path / filepath
if not await asyncio.to_thread(file_path.exists):
return None
async with aiofiles.open(file_path, 'rb') as f:
content = await f.read()
index = await self._load_index()
meta = index.get('media', {}).get(filepath, {})
return content, meta
async def finish_crawl(self, stats: Dict[str, Any]) -> Optional[str]:
"""Finish crawl and commit if changes."""
media_new = stats.get('media_new', 0)
if media_new > 0:
return await self._git.commit(f"Crawl: {media_new} media files")
return None
async def get_stats(self) -> Dict[str, Any]:
"""Get vault statistics."""
index = await self._load_index()
commits = await self._git.log(5)
return {
'domain': self.domain,
'path': str(self.path),
'media_count': len(index.get('media', {})),
'recent_commits': commits,
}
class DomainLinkpeekVault:
"""
Git LFS-tracked screenshot vault for a single domain.
Stores page screenshots in a tree structure mirroring the original domain.
One screenshot per page, named after the URL path.
"""
def __init__(self, vault_base: str, domain: str):
self.domain = domain.lower()
self.base_path = Path(vault_base)
self.path = get_filevault_path(self.base_path, self.domain)
self.index_path = self.path / 'index.json'
self._git = GitRepo(self.path, use_lfs=True)
self._initialized = False
async def init(self) -> bool:
"""Initialize vault. Returns True if newly created."""
is_new = await self._git.init()
if not await asyncio.to_thread(self.index_path.exists):
await self._save_index({'screenshots': {}, 'created': datetime.now(timezone.utc).isoformat()})
is_new = True
self._initialized = True
return is_new
async def _load_index(self) -> Dict[str, Any]:
if not await asyncio.to_thread(self.index_path.exists):
return {'screenshots': {}}
async with aiofiles.open(self.index_path, 'r') as f:
return json.loads(await f.read())
async def _save_index(self, index: Dict[str, Any]):
async with aiofiles.open(self.index_path, 'w') as f:
await f.write(json.dumps(index, indent=2))
def url_to_screenshot_path(self, url: str) -> str:
"""
Convert URL to screenshot filepath.
Examples:
https://example.com/ -> index.png
https://example.com/about -> about/index.png
https://example.com/blog/post.html -> blog/post.png
"""
parsed = urlparse(url)
path = parsed.path.strip('/')
if not path:
return 'index.png'
# Replace extension with .png or add /index.png for directories
if '.' in path.split('/')[-1]:
# Has extension, replace with .png
base = path.rsplit('.', 1)[0]
return f"{base}.png"
else:
# Directory-style URL
return f"{path}/index.png"
async def archive_screenshot(
self,
url: str,
screenshot_data: bytes,
) -> Tuple[bool, str, Path]:
"""
Archive a page screenshot.
Returns (is_new, content_hash, file_path).
"""
if not self._initialized:
await self.init()
filepath = self.url_to_screenshot_path(url)
file_path = self.path / filepath
await asyncio.to_thread(file_path.parent.mkdir, parents=True, exist_ok=True)
# Check if new or changed
chash = content_hash(screenshot_data)
exists = await asyncio.to_thread(file_path.exists)
is_new = not exists
if exists:
async with aiofiles.open(file_path, 'rb') as f:
existing_hash = content_hash(await f.read())
is_new = existing_hash != chash
# Save screenshot
async with aiofiles.open(file_path, 'wb') as f:
await f.write(screenshot_data)
# Update index
index = await self._load_index()
index['screenshots'][filepath] = {
'url': url,
'hash': chash,
'size': len(screenshot_data),
'updated': datetime.now(timezone.utc).isoformat(),
}
await self._save_index(index)
return is_new, chash, file_path
async def get_screenshot(self, url: str) -> Optional[Tuple[bytes, Dict[str, Any]]]:
"""Get archived screenshot and metadata."""
filepath = self.url_to_screenshot_path(url)
file_path = self.path / filepath
if not await asyncio.to_thread(file_path.exists):
return None
async with aiofiles.open(file_path, 'rb') as f:
content = await f.read()
index = await self._load_index()
meta = index.get('screenshots', {}).get(filepath, {})
return content, meta
async def finish_crawl(self, stats: Dict[str, Any]) -> Optional[str]:
"""Finish crawl and commit if changes."""
screenshots_new = stats.get('screenshots_new', 0)
if screenshots_new > 0:
return await self._git.commit(f"Crawl: {screenshots_new} screenshots")
return None
async def get_stats(self) -> Dict[str, Any]:
"""Get vault statistics."""
index = await self._load_index()
commits = await self._git.log(5)
return {
'domain': self.domain,
'path': str(self.path),
'screenshots_count': len(index.get('screenshots', {})),
'recent_commits': commits,
}
class VaultManager:
"""
Manages HTML, Media, and Linkpeek (screenshot) vaults for multiple domains.
"""
def __init__(
self,
html_vault_base: str = 'html_vault',
media_vault_base: str = 'media_vault',
linkpeek_vault_base: str = 'linkpeek_vault',
media_base_url: str = '/media',
linkpeek_base_url: str = '/linkpeek',
):
self.html_base = Path(html_vault_base)
self.media_base = Path(media_vault_base)
self.linkpeek_base = Path(linkpeek_vault_base)
self.media_base_url = media_base_url
self.linkpeek_base_url = linkpeek_base_url
self._html_vaults: Dict[str, DomainHtmlVault] = {}
self._media_vaults: Dict[str, DomainMediaVault] = {}
self._linkpeek_vaults: Dict[str, DomainLinkpeekVault] = {}
def get_html_vault(self, domain: str) -> DomainHtmlVault:
"""Get or create HTML vault for domain."""
domain = domain.lower()
if domain not in self._html_vaults:
self._html_vaults[domain] = DomainHtmlVault(
str(self.html_base), domain, self.media_base_url
)
return self._html_vaults[domain]
def get_media_vault(self, domain: str) -> DomainMediaVault:
"""Get or create media vault for domain."""
domain = domain.lower()
if domain not in self._media_vaults:
self._media_vaults[domain] = DomainMediaVault(str(self.media_base), domain)
return self._media_vaults[domain]
def get_linkpeek_vault(self, domain: str) -> DomainLinkpeekVault:
"""Get or create linkpeek (screenshot) vault for domain."""
domain = domain.lower()
if domain not in self._linkpeek_vaults:
self._linkpeek_vaults[domain] = DomainLinkpeekVault(str(self.linkpeek_base), domain)
return self._linkpeek_vaults[domain]
def get_vaults_for_url(self, url: str) -> Tuple[DomainHtmlVault, DomainMediaVault, DomainLinkpeekVault]:
"""Get HTML, media, and linkpeek vaults for URL's domain."""
parsed = urlparse(url)
domain = parsed.netloc or parsed.path.split('/')[0]
return self.get_html_vault(domain), self.get_media_vault(domain), self.get_linkpeek_vault(domain)
def get_linkpeek_url(self, url: str) -> str:
"""
Get linkpeek URL for a page screenshot.
Example: https://example.com/about
-> /linkpeek/{9-layers}/{salted_hash}/about/index.png
"""
parsed = urlparse(url)
domain = parsed.netloc.lower()
# Get screenshot path
vault = self.get_linkpeek_vault(domain)
screenshot_path = vault.url_to_screenshot_path(url)
# Get salted hash for domain (same across all vaults)
h = domain_hash(domain, salted=True)
layers = '/'.join([h[i:i+2] for i in range(0, 18, 2)])
return f"{self.linkpeek_base_url}/{layers}/{h}/{screenshot_path}"
async def list_domains(self, vault_type: str = 'html') -> List[str]:
"""List all domains with vaults."""
if vault_type == 'html':
base = self.html_base
elif vault_type == 'media':
base = self.media_base
else: # linkpeek
base = self.linkpeek_base
def _walk_domains():
domains = []
if not base.exists():
return domains
for root, dirs, files in os.walk(base):
if '.git' in dirs:
rel_path = Path(root).relative_to(base)
parts = list(rel_path.parts)
if len(parts) >= 10: # 9 hash layers + domain
domains.append(parts[-1])
return sorted(set(domains))
return await asyncio.to_thread(_walk_domains)
async def get_all_stats(self) -> Dict[str, List[Dict[str, Any]]]:
"""Get stats for all domains."""
html_domains = await self.list_domains('html')
media_domains = await self.list_domains('media')
linkpeek_domains = await self.list_domains('linkpeek')
html_stats = []
for domain in html_domains:
vault = self.get_html_vault(domain)
await vault.init()
html_stats.append(await vault.get_stats())
media_stats = []
for domain in media_domains:
vault = self.get_media_vault(domain)
await vault.init()
media_stats.append(await vault.get_stats())
linkpeek_stats = []
for domain in linkpeek_domains:
vault = self.get_linkpeek_vault(domain)
await vault.init()
linkpeek_stats.append(await vault.get_stats())
return {
'html_vaults': html_stats,
'media_vaults': media_stats,
'linkpeek_vaults': linkpeek_stats,
}

545
filevault.py Normal file
View file

@ -0,0 +1,545 @@
"""
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)

264
neopig.py
View file

@ -27,6 +27,8 @@ from datetime import datetime, timezone
from pathlib import Path
from typing import List, Dict, Any, Optional, Set
from urllib.parse import urlparse
from async_web_fetcher import (
AsyncWebFetcher,
CrawlMode,
@ -37,6 +39,8 @@ from async_web_fetcher import (
)
from storage import ImageVault
from database import Database
from screenshot import ScreenshotCapture, ScreenshotConfig
from domain_vault import VaultManager, DomainHtmlVault, DomainMediaVault, DomainLinkpeekVault, extract_media_urls
logging.basicConfig(
level=logging.INFO,
@ -54,23 +58,112 @@ class NeoPig:
self,
db_path: str = "neopig.db",
vault_path: str = "vault",
user_agent: str = "neopig/1.0 (ethical image crawler)"
user_agent: str = "neopig/1.0 (ethical image crawler)",
screenshot_config: ScreenshotConfig = None,
):
self.db = Database(db_path)
self.vault = ImageVault(vault_path)
# Triple filevault system: html_vault/, media_vault/, and linkpeek_vault/
self.domain_vaults = VaultManager(
html_vault_base=f"{vault_path}/html_vault",
media_vault_base=f"{vault_path}/media_vault",
linkpeek_vault_base=f"{vault_path}/linkpeek_vault",
media_base_url='/media',
linkpeek_base_url='/linkpeek',
)
self.fetcher = AsyncWebFetcher(user_agent=user_agent)
self.screenshot = ScreenshotCapture(screenshot_config or ScreenshotConfig())
self.screenshot_config = screenshot_config or ScreenshotConfig()
self.vault_path = vault_path
# Track stats
self.stats = {
'pages_crawled': 0,
'pages_changed': 0, # Pages with content changes (for git commit)
'media_found': 0,
'media_downloaded': 0,
'media_new': 0, # New media (for git commit)
'duplicates_skipped': 0,
'screenshots_taken': 0,
'errors': 0,
}
# Track seen media URLs to avoid re-processing
self.seen_media: Set[str] = set()
# Track screenshotted pages to avoid duplicates
self.seen_screenshots: Set[str] = set()
# Track per-domain stats for vault commits
self._domain_stats: Dict[str, Dict[str, int]] = {} # domain -> {pages_changed, media_new, screenshots_new}
def _get_domain(self, url: str) -> str:
"""Extract domain from URL."""
parsed = urlparse(url)
return parsed.netloc.lower()
def _track_domain_stat(self, domain: str, stat: str, increment: int = 1):
"""Track per-domain stats for vault commits."""
if domain not in self._domain_stats:
self._domain_stats[domain] = {'pages_changed': 0, 'media_new': 0, 'screenshots_new': 0}
self._domain_stats[domain][stat] += increment
async def _archive_page_to_vault(
self,
url: str,
html: str,
media_mappings: Dict[str, str] = None,
):
"""Archive a page to the HTML vault."""
domain = self._get_domain(url)
html_vault = self.domain_vaults.get_html_vault(domain)
is_changed, _ = await html_vault.archive_page(url, html, media_mappings)
if is_changed:
self._track_domain_stat(domain, 'pages_changed')
self.stats['pages_changed'] += 1
async def _archive_media_to_vault(
self,
url: str,
content: bytes,
page_url: str = '',
):
"""Archive media to the media vault."""
domain = self._get_domain(url)
media_vault = self.domain_vaults.get_media_vault(domain)
is_new, _, _ = await media_vault.archive_media(url, content, page_url)
if is_new:
self._track_domain_stat(domain, 'media_new')
self.stats['media_new'] += 1
async def _archive_screenshot_to_vault(
self,
url: str,
screenshot_data: bytes,
):
"""Archive screenshot to the linkpeek vault."""
domain = self._get_domain(url)
linkpeek_vault = self.domain_vaults.get_linkpeek_vault(domain)
is_new, _, _ = await linkpeek_vault.archive_screenshot(url, screenshot_data)
if is_new:
self._track_domain_stat(domain, 'screenshots_new')
async def _finish_domain_vaults(self, keywords: List[str] = None):
"""Commit changes to all domain vaults that have diffs."""
for domain, stats in self._domain_stats.items():
if stats['pages_changed'] > 0 or stats['media_new'] > 0 or stats['screenshots_new'] > 0:
html_vault = self.domain_vaults.get_html_vault(domain)
media_vault = self.domain_vaults.get_media_vault(domain)
linkpeek_vault = self.domain_vaults.get_linkpeek_vault(domain)
html_commit = await html_vault.finish_crawl(stats)
media_commit = await media_vault.finish_crawl(stats)
linkpeek_commit = await linkpeek_vault.finish_crawl(stats)
if html_commit:
logger.info(f"Committed HTML vault for {domain}: {html_commit[:8]}")
if media_commit:
logger.info(f"Committed media vault for {domain}: {media_commit[:8]}")
if linkpeek_commit:
logger.info(f"Committed linkpeek vault for {domain}: {linkpeek_commit[:8]}")
async def init(self):
"""Initialize database and vault."""
@ -115,6 +208,32 @@ class NeoPig:
logger.info(f"Keywords: {keywords}")
logger.info(f"Depth: {'unlimited' if depth == -1 else depth}")
# Track timing for stats
start_time = datetime.now(timezone.utc)
last_stats_time = start_time
stats_running = True
# Background task to emit stats every 15 seconds
async def stats_reporter():
nonlocal last_stats_time
while stats_running:
await asyncio.sleep(15)
if not stats_running:
break
elapsed = (datetime.now(timezone.utc) - start_time).total_seconds()
rate = self.stats['media_downloaded'] / elapsed * 60 if elapsed > 0 else 0
logger.info(f"=== CRAWL STATS ({elapsed:.0f}s) ===")
logger.info(f" Pages: {self.stats['pages_crawled']} | "
f"Found: {self.stats['media_found']} | "
f"Downloaded: {self.stats['media_downloaded']} | "
f"Dupes: {self.stats['duplicates_skipped']} | "
f"Screenshots: {self.stats['screenshots_taken']} | "
f"Errors: {self.stats['errors']}")
logger.info(f" Rate: {rate:.1f}/min | "
f"Vault size: {len(self.seen_media)}")
stats_task = asyncio.create_task(stats_reporter())
# Media callback - called for each discovered media item
async def on_media_discovered(item: Dict[str, Any]):
url = item.get('url')
@ -126,6 +245,8 @@ class NeoPig:
if download_media:
await self._process_media_item(item, job_id, keywords)
# Note: Screenshots are now captured per-page in on_page_fetched,
# not per-media-item, to honor crawl delay as a unit
# Progress callback
async def on_progress(msg: str):
@ -135,6 +256,17 @@ class NeoPig:
f"{self.stats['media_found']} media found, "
f"{self.stats['media_downloaded']} downloaded")
# Page callback - archive raw HTML to vault and capture screenshot
# Screenshot happens here (same crawl delay window as page fetch)
async def on_page_fetched(url: str, html: str):
# For now, archive without media URL rewriting (we'd need to download media first)
# TODO: Build media_mappings after media is downloaded
await self._archive_page_to_vault(url, html, media_mappings=None)
# Capture screenshot for every page (honors crawl delay as a UNIT with page fetch)
if self.screenshot_config.enabled:
await self._capture_page_screenshot(url, job_id, page_title='')
# Run the crawl
pages = await self.fetcher.fetch_with_depth(
start_url=target_uri,
@ -144,19 +276,37 @@ class NeoPig:
mode=mode,
media_callback=on_media_discovered,
progress_callback=on_progress,
page_callback=on_page_fetched,
)
self.stats['pages_crawled'] = len(pages)
# Stop the stats reporter
stats_running = False
stats_task.cancel()
try:
await stats_task
except asyncio.CancelledError:
pass
# Calculate final stats
elapsed = (datetime.now(timezone.utc) - start_time).total_seconds()
rate = self.stats['media_downloaded'] / elapsed * 60 if elapsed > 0 else 0
# Update job status
await self.db.complete_crawl_job(job_id, self.stats)
logger.info(f"Crawl complete!")
logger.info(f" Pages crawled: {self.stats['pages_crawled']}")
logger.info(f" Media found: {self.stats['media_found']}")
logger.info(f" Media downloaded: {self.stats['media_downloaded']}")
logger.info(f" Duplicates skipped: {self.stats['duplicates_skipped']}")
logger.info(f" Errors: {self.stats['errors']}")
# Commit domain vaults if there were changes
await self._finish_domain_vaults(keywords)
logger.info(f"=== CRAWL COMPLETE ({elapsed:.0f}s) ===")
logger.info(f" Pages: {self.stats['pages_crawled']} | "
f"Found: {self.stats['media_found']} | "
f"Downloaded: {self.stats['media_downloaded']} | "
f"Dupes: {self.stats['duplicates_skipped']} | "
f"Screenshots: {self.stats['screenshots_taken']} | "
f"Errors: {self.stats['errors']}")
logger.info(f" Rate: {rate:.1f}/min | Total time: {elapsed:.1f}s")
return self.stats
@ -223,6 +373,9 @@ class NeoPig:
ext = self._get_extension(media_uri, result.get('mime_type', ''))
await self.vault.store(md5_hash, result['data'], ext)
# Archive to domain media vault (git-tracked)
await self._archive_media_to_vault(media_uri, result['data'], page_uri)
# Record in database with full context
await self.db.create_media_record(
md5_hash=md5_hash,
@ -246,6 +399,58 @@ class NeoPig:
logger.warning(f"Failed to process {media_uri}: {e}")
self.stats['errors'] += 1
async def _capture_page_screenshot(
self,
page_uri: str,
job_id: int,
page_title: str = '',
):
"""Capture and store a screenshot of a page."""
if not self.screenshot_config.enabled:
return
if page_uri in self.seen_screenshots:
return
self.seen_screenshots.add(page_uri)
try:
result = await self.screenshot.capture(page_uri)
if not result:
return
md5_hash = result['md5_hash']
screenshot_data = result['data']
# Store in MD5 vault (for deduplication)
if not await self.vault.exists(md5_hash):
await self.vault.store(md5_hash, screenshot_data, 'png')
# Archive to linkpeek vault (git-tracked by URL path)
await self._archive_screenshot_to_vault(page_uri, screenshot_data)
# Record in database as screenshot type
await self.db.create_media_record(
md5_hash=md5_hash,
media_uri=f"screenshot:{page_uri}",
page_uri=page_uri,
crawl_job_id=job_id,
media_type='screenshot',
mime_type='image/png',
file_size=result.get('size', 0),
page_title=page_title,
page_description='',
page_keywords='',
alt_text=f"Screenshot of {page_uri}",
link_text='',
)
self.stats['screenshots_taken'] += 1
logger.debug(f"Screenshot captured: {page_uri} -> {md5_hash}")
except Exception as e:
logger.warning(f"Screenshot failed for {page_uri}: {e}")
def _get_extension(self, url: str, mime_type: str) -> str:
"""Determine file extension from URL or MIME type."""
from urllib.parse import urlparse
@ -338,6 +543,34 @@ async def main():
help="Verbose output"
)
# Screenshot options (all off by default)
parser.add_argument(
"--screenshot",
action="store_true",
help="Enable page screenshots (requires uri2png)"
)
parser.add_argument(
"--screenshot-width",
type=int,
default=1280,
help="Screenshot viewport width in pixels (default: 1280)"
)
parser.add_argument(
"--screenshot-height",
type=int,
default=1024,
help="Screenshot viewport height in pixels (default: 1024)"
)
parser.add_argument(
"--screenshot-delay",
type=int,
default=1000,
help="Delay after page load in ms (default: 1000)"
)
args = parser.parse_args()
if args.verbose:
@ -353,8 +586,23 @@ async def main():
}
mode = mode_map[args.mode]
# Create screenshot config from args
screenshot_config = ScreenshotConfig(
enabled=args.screenshot,
width=args.screenshot_width,
height=args.screenshot_height,
delay=args.screenshot_delay,
)
if args.screenshot:
logger.info(f"Screenshots enabled: {screenshot_config.width}x{screenshot_config.height}, delay={screenshot_config.delay}ms")
# Initialize and run
pig = NeoPig(db_path=args.db, vault_path=args.vault)
pig = NeoPig(
db_path=args.db,
vault_path=args.vault,
screenshot_config=screenshot_config,
)
await pig.init()
# Load previously crawled media URIs to enable resume

182
screenshot.py Normal file
View file

@ -0,0 +1,182 @@
#!/usr/bin/env python3
"""
Screenshot capture module for neopig.
Wraps uri2png for async-compatible page screenshots.
Screenshots are stored in vault with MD5 hash like other media.
"""
import asyncio
import hashlib
import logging
import subprocess
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import Optional
logger = logging.getLogger(__name__)
@dataclass
class ScreenshotConfig:
"""Screenshot capture configuration."""
enabled: bool = False
width: int = 1280
height: int = 1024
delay: int = 1000 # ms after DOM load
user_agent: Optional[str] = None
class ScreenshotCapture:
"""
Async screenshot capture using uri2png.
Since uri2png uses GTK main loop, we run it in a subprocess
to avoid blocking the async event loop.
"""
def __init__(self, config: ScreenshotConfig = None):
self.config = config or ScreenshotConfig()
self._uri2png_available = None
async def is_available(self) -> bool:
"""Check if uri2png is installed and available."""
if self._uri2png_available is not None:
return self._uri2png_available
try:
proc = await asyncio.create_subprocess_exec(
'python', '-c', 'from uri2png import Uri2Png',
stdout=asyncio.subprocess.DEVNULL,
stderr=asyncio.subprocess.DEVNULL,
)
await proc.wait()
self._uri2png_available = proc.returncode == 0
except Exception:
self._uri2png_available = False
if not self._uri2png_available:
logger.warning("uri2png not available - screenshots disabled")
return self._uri2png_available
async def capture(self, uri: str) -> Optional[dict]:
"""
Capture screenshot of a URI.
Returns:
Dict with 'data', 'md5_hash', 'mime_type' or None on failure
"""
if not self.config.enabled:
return None
if not await self.is_available():
return None
# Create temp file for screenshot
with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp:
tmp_path = tmp.name
try:
# Build uri2png command
cmd = [
'python', '-c',
f'''
from uri2png import Uri2Png
Uri2Png(
uri="{uri}",
filepath="{tmp_path}",
width={self.config.width},
height={self.config.height},
delay={self.config.delay},
user_agent={repr(self.config.user_agent)},
).capture()
'''
]
# Run with timeout (uri2png can hang on bad URLs)
proc = await asyncio.create_subprocess_exec(
*cmd,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
)
try:
stdout, stderr = await asyncio.wait_for(
proc.communicate(),
timeout=30.0 # 30 second timeout
)
except asyncio.TimeoutError:
proc.kill()
await proc.wait()
logger.warning(f"Screenshot timeout: {uri}")
return None
if proc.returncode != 0:
logger.warning(f"Screenshot failed ({proc.returncode}): {uri}")
if stderr:
logger.debug(f"stderr: {stderr.decode()}")
return None
# Read screenshot data
path = Path(tmp_path)
def _check_file():
if not path.exists():
return None
size = path.stat().st_size
if size == 0:
return None
return path.read_bytes()
data = await asyncio.to_thread(_check_file)
if data is None:
logger.warning(f"Screenshot empty: {uri}")
return None
md5_hash = hashlib.md5(data).hexdigest()
logger.debug(f"Screenshot captured: {uri} -> {md5_hash}")
return {
'data': data,
'md5_hash': md5_hash,
'mime_type': 'image/png',
'size': len(data),
'source_uri': uri,
}
except Exception as e:
logger.warning(f"Screenshot error for {uri}: {e}")
return None
finally:
# Cleanup temp file (sync unlink is fine in finally - small operation)
def _cleanup():
try:
Path(tmp_path).unlink(missing_ok=True)
except Exception:
pass
try:
await asyncio.to_thread(_cleanup)
except Exception:
pass
async def capture_to_file(self, uri: str, output_path: str) -> bool:
"""
Capture screenshot directly to a file.
Returns:
True on success, False on failure
"""
result = await self.capture(uri)
if result is None:
return False
try:
await asyncio.to_thread(Path(output_path).write_bytes, result['data'])
return True
except Exception as e:
logger.warning(f"Failed to write screenshot: {e}")
return False

258
serp.py
View file

@ -34,7 +34,15 @@ import uvicorn
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
app = FastAPI(title="neopig SERP", description="Search hydrated media")
app = FastAPI(title="neopig", description="Media crawler + SERP + Screenshot service")
# Try to include uri2png screenshot router (optional dependency)
try:
from uri2png import get_screenshot_router
app.include_router(get_screenshot_router())
logger.info("Screenshot router loaded from uri2png")
except ImportError:
logger.warning("uri2png not installed - screenshot endpoints not available")
# Config - set via startup
DB_PATH = "neopig.db"
@ -44,9 +52,83 @@ VAULT_PATH = Path("vault")
ACTIVE_CRAWLS: Dict[int, Dict[str, Any]] = {}
async def init_database():
"""Initialize database schema if needed."""
async with aiosqlite.connect(DB_PATH) as db:
# Crawl jobs table
await db.execute("""
CREATE TABLE IF NOT EXISTS crawl_jobs (
id INTEGER PRIMARY KEY AUTOINCREMENT,
target_uri TEXT NOT NULL,
keywords TEXT,
mode TEXT DEFAULT 'images',
status TEXT DEFAULT 'running',
started_at TEXT NOT NULL,
completed_at TEXT,
stats TEXT
)
""")
# Media records table
await db.execute("""
CREATE TABLE IF NOT EXISTS media (
md5_hash TEXT PRIMARY KEY,
media_type TEXT,
mime_type TEXT,
file_size INTEGER,
keywords TEXT,
alt_text TEXT,
title TEXT,
first_seen_at TEXT NOT NULL,
analysis_status TEXT DEFAULT 'pending',
analysis_result TEXT
)
""")
# Media sources table
await db.execute("""
CREATE TABLE IF NOT EXISTS media_sources (
id INTEGER PRIMARY KEY AUTOINCREMENT,
md5_hash TEXT NOT NULL,
media_uri TEXT NOT NULL,
page_uri TEXT,
page_title TEXT,
page_description TEXT,
page_keywords TEXT,
alt_text TEXT,
link_text TEXT,
crawl_job_id INTEGER,
discovered_at TEXT NOT NULL,
FOREIGN KEY (md5_hash) REFERENCES media(md5_hash),
FOREIGN KEY (crawl_job_id) REFERENCES crawl_jobs(id),
UNIQUE(md5_hash, media_uri, page_uri)
)
""")
# Indexes
await db.execute("CREATE INDEX IF NOT EXISTS idx_media_type ON media(media_type)")
await db.execute("CREATE INDEX IF NOT EXISTS idx_media_analysis ON media(analysis_status)")
await db.execute("CREATE INDEX IF NOT EXISTS idx_sources_hash ON media_sources(md5_hash)")
await db.execute("CREATE INDEX IF NOT EXISTS idx_sources_job ON media_sources(crawl_job_id)")
await db.execute("CREATE INDEX IF NOT EXISTS idx_sources_media_uri ON media_sources(media_uri)")
await db.execute("CREATE INDEX IF NOT EXISTS idx_sources_page_uri ON media_sources(page_uri)")
await db.commit()
logger.info(f"Database initialized: {DB_PATH}")
@app.on_event("startup")
async def startup_event():
"""Initialize database on startup."""
await init_database()
VAULT_PATH.mkdir(parents=True, exist_ok=True)
logger.info(f"Vault directory ready: {VAULT_PATH}")
class CrawlRequest(BaseModel):
"""Request to start a new crawl."""
target_uri: str
targets: List[str] = [] # Multiple target URIs
target_uri: str = "" # Deprecated: single target (for backwards compat)
keywords: List[str] = []
mode: str = "images" # text, images, videos, media, all
depth: int = -1 # -1 = unlimited
@ -459,8 +541,8 @@ CRAWL_HTML = """
<form id="crawl-form" onsubmit="startCrawl(event)">
<div class="form-group">
<label>Target URI</label>
<input type="text" id="target" placeholder="https://example.com" required>
<label>Target URIs (space, comma, or newline separated)</label>
<textarea id="target" placeholder="https://example.com https://another.com" rows="2" required style="width:100%;padding:10px;border:1px solid #333;border-radius:4px;background:#1a1a1a;color:#e0e0e0;font-size:14px;resize:vertical;"></textarea>
</div>
<div class="form-group">
@ -480,12 +562,12 @@ CRAWL_HTML = """
</select>
</div>
<div class="form-group">
<label>Depth (-1 = unlimited)</label>
<input type="number" id="depth" value="-1" min="-1">
<label>Depth (5 recommended, max 15)</label>
<input type="number" id="depth" value="5" min="1" max="15">
</div>
<div class="form-group">
<label>Max Pages (-1 = unlimited)</label>
<input type="number" id="max_pages" value="-1" min="-1">
<input type="number" id="max_pages" value="500" min="-1">
</div>
</div>
@ -516,8 +598,22 @@ CRAWL_HTML = """
.map(k => k.trim())
.filter(k => k.length > 0);
// Parse multiple target URIs (space, comma, or newline separated)
const targetsRaw = document.getElementById('target').value;
const targets = targetsRaw
.split(/[,\\s\\n]+/)
.map(t => t.trim())
.filter(t => t.length > 0 && t.startsWith('http'));
if (targets.length === 0) {
alert('Please enter at least one valid URI (must start with http)');
btn.disabled = false;
btn.textContent = 'Start Crawl';
return;
}
const payload = {
target_uri: document.getElementById('target').value,
targets: targets,
keywords: keywords,
mode: document.getElementById('mode').value,
depth: parseInt(document.getElementById('depth').value),
@ -536,8 +632,9 @@ CRAWL_HTML = """
const err = await res.json();
alert('Error: ' + (err.detail || 'Failed to start crawl'));
} else {
const job = await res.json();
alert('Crawl started! Job ID: ' + job.job_id);
const result = await res.json();
const jobIds = result.job_ids || [result.job_id];
alert('Crawl started! Job IDs: ' + jobIds.join(', '));
loadJobs();
}
} catch (err) {
@ -966,6 +1063,26 @@ async def view_media_page(md5_hash: str):
"""
@app.get("/health")
async def health():
"""Health check endpoint."""
has_screenshot = False
try:
from uri2png import get_available_engines
has_screenshot = True
except ImportError:
pass
return {
"status": "healthy",
"features": {
"search": True,
"crawl": True,
"screenshot": has_screenshot
}
}
@app.get("/api/stats")
async def get_stats():
"""Get database statistics."""
@ -1158,9 +1275,10 @@ async def get_crawl_job(job_id: int):
@app.post("/api/crawl")
async def start_crawl(request: CrawlRequest, background_tasks: BackgroundTasks):
"""
Start a new crawl job.
Start new crawl job(s).
The crawl runs in the background. Poll /api/crawl/jobs/{id} for status.
Supports multiple targets - creates one job per target.
The crawls run in the background. Poll /api/crawl/jobs/{id} for status.
"""
# Import neopig here to avoid circular imports
from neopig import NeoPig
@ -1176,65 +1294,79 @@ async def start_crawl(request: CrawlRequest, background_tasks: BackgroundTasks):
}
mode = mode_map.get(request.mode, CrawlMode.IMAGES)
# Create job record first
async with aiosqlite.connect(DB_PATH) as db:
cursor = await db.execute(
"""
INSERT INTO crawl_jobs (target_uri, keywords, mode, status, started_at)
VALUES (?, ?, ?, 'running', datetime('now'))
""",
(request.target_uri, json.dumps(request.keywords), request.mode)
)
await db.commit()
job_id = cursor.lastrowid
# Get targets (support both new 'targets' array and old 'target_uri' single value)
targets = request.targets if request.targets else [request.target_uri] if request.target_uri else []
if not targets:
raise HTTPException(status_code=400, detail="No target URIs provided")
# Run crawl in background
async def run_crawl():
try:
pig = NeoPig(db_path=DB_PATH, vault_path=str(VAULT_PATH))
await pig.init()
job_ids = []
# Override job_id since we already created it
pig.db = None # We'll update directly
stats = await pig.crawl(
target_uri=request.target_uri,
keywords=request.keywords,
mode=mode,
depth=request.depth,
max_pages=request.max_pages,
download_media=request.download_media,
# Create a job for each target
for target_uri in targets:
async with aiosqlite.connect(DB_PATH) as db:
cursor = await db.execute(
"""
INSERT INTO crawl_jobs (target_uri, keywords, mode, status, started_at)
VALUES (?, ?, ?, 'running', datetime('now'))
""",
(target_uri, json.dumps(request.keywords), request.mode)
)
await db.commit()
job_id = cursor.lastrowid
job_ids.append(job_id)
# Update job as completed
async with aiosqlite.connect(DB_PATH) as db:
await db.execute(
"""
UPDATE crawl_jobs
SET status = 'completed', completed_at = datetime('now'), stats = ?
WHERE id = ?
""",
(json.dumps(stats), job_id)
# Run crawl in background (closure captures job_id and target_uri)
async def run_crawl(jid=job_id, uri=target_uri):
try:
pig = NeoPig(db_path=DB_PATH, vault_path=str(VAULT_PATH))
await pig.init()
# Load existing seen media for resume capability
crawled_media = await pig.db.get_crawled_media_uris()
if crawled_media:
pig.seen_media = crawled_media
# Cap depth at 15 (convert -1 or values > 15 to 15)
depth = request.depth if 1 <= request.depth <= 15 else 15
stats = await pig.crawl(
target_uri=uri,
keywords=request.keywords,
mode=mode,
depth=depth,
max_pages=request.max_pages,
download_media=request.download_media,
)
await db.commit()
except Exception as e:
logger.error(f"Crawl job {job_id} failed: {e}")
async with aiosqlite.connect(DB_PATH) as db:
await db.execute(
"""
UPDATE crawl_jobs
SET status = 'failed', completed_at = datetime('now'), stats = ?
WHERE id = ?
""",
(json.dumps({"error": str(e)}), job_id)
)
await db.commit()
# Update job as completed
async with aiosqlite.connect(DB_PATH) as db:
await db.execute(
"""
UPDATE crawl_jobs
SET status = 'completed', completed_at = datetime('now'), stats = ?
WHERE id = ?
""",
(json.dumps(stats), jid)
)
await db.commit()
# Schedule background task
background_tasks.add_task(asyncio.create_task, run_crawl())
except Exception as e:
logger.error(f"Crawl job {jid} failed: {e}")
async with aiosqlite.connect(DB_PATH) as db:
await db.execute(
"""
UPDATE crawl_jobs
SET status = 'failed', completed_at = datetime('now'), stats = ?
WHERE id = ?
""",
(json.dumps({"error": str(e)}), jid)
)
await db.commit()
return {"job_id": job_id, "status": "running"}
# Schedule background task
background_tasks.add_task(asyncio.create_task, run_crawl())
return {"job_ids": job_ids, "status": "running", "count": len(job_ids)}
def main():

138
storage.py
View file

@ -50,11 +50,7 @@ class ImageVault:
return
if HAS_FILEVAULT:
loop = asyncio.get_event_loop()
self.filevault = await loop.run_in_executor(
None,
lambda: FileVault(str(self.vault_path))
)
self.filevault = await asyncio.to_thread(FileVault, str(self.vault_path))
else:
# Fallback: create directory structure
await aiofiles.os.makedirs(self.vault_path, exist_ok=True)
@ -79,12 +75,7 @@ class ImageVault:
"""
if HAS_FILEVAULT and self.filevault:
# filevault stores by content hash automatically
loop = asyncio.get_event_loop()
key = await loop.run_in_executor(
None,
lambda: self.filevault.put(data)
)
return key
return await asyncio.to_thread(self.filevault.put, data)
else:
# Fallback: store in subdirectory by first 2 chars
subdir = md5_hash[:2]
@ -99,102 +90,109 @@ class ImageVault:
async def exists(self, md5_hash: str) -> bool:
"""Check if a file exists by MD5 hash."""
if HAS_FILEVAULT and self.filevault:
loop = asyncio.get_event_loop()
return await loop.run_in_executor(
None,
lambda: self.filevault.exists(md5_hash)
)
return await asyncio.to_thread(self.filevault.exists, md5_hash)
else:
# Check for any file starting with this hash
subdir = self.vault_path / md5_hash[:2]
if not subdir.exists():
def _check_exists():
if not subdir.exists():
return False
for f in subdir.iterdir():
if f.name.startswith(md5_hash):
return True
return False
for f in subdir.iterdir():
if f.name.startswith(md5_hash):
return True
return False
return await asyncio.to_thread(_check_exists)
async def get(self, md5_hash: str) -> Optional[bytes]:
"""Retrieve data by MD5 hash."""
if HAS_FILEVAULT and self.filevault:
loop = asyncio.get_event_loop()
return await loop.run_in_executor(
None,
lambda: self.filevault.get(md5_hash)
)
return await asyncio.to_thread(self.filevault.get, md5_hash)
else:
subdir = self.vault_path / md5_hash[:2]
if not subdir.exists():
def _find_file():
if not subdir.exists():
return None
for f in subdir.iterdir():
if f.name.startswith(md5_hash):
return f
return None
for f in subdir.iterdir():
if f.name.startswith(md5_hash):
async with aiofiles.open(f, 'rb') as file:
return await file.read()
return None
file_path = await asyncio.to_thread(_find_file)
if file_path is None:
return None
async with aiofiles.open(file_path, 'rb') as file:
return await file.read()
async def get_path(self, md5_hash: str) -> Optional[Path]:
"""Get the filesystem path for a stored file."""
if HAS_FILEVAULT and self.filevault:
# filevault may not expose paths directly
loop = asyncio.get_event_loop()
path = await loop.run_in_executor(
None,
lambda: self.filevault.path(md5_hash) if hasattr(self.filevault, 'path') else None
)
def _get_path():
return self.filevault.path(md5_hash) if hasattr(self.filevault, 'path') else None
path = await asyncio.to_thread(_get_path)
return Path(path) if path else None
else:
subdir = self.vault_path / md5_hash[:2]
if not subdir.exists():
def _find_file():
if not subdir.exists():
return None
for f in subdir.iterdir():
if f.name.startswith(md5_hash):
return f
return None
for f in subdir.iterdir():
if f.name.startswith(md5_hash):
return f
return None
return await asyncio.to_thread(_find_file)
async def delete(self, md5_hash: str) -> bool:
"""Delete a file by MD5 hash."""
if HAS_FILEVAULT and self.filevault:
loop = asyncio.get_event_loop()
return await loop.run_in_executor(
None,
lambda: self.filevault.delete(md5_hash)
)
return await asyncio.to_thread(self.filevault.delete, md5_hash)
else:
subdir = self.vault_path / md5_hash[:2]
if not subdir.exists():
def _find_file():
if not subdir.exists():
return None
for f in subdir.iterdir():
if f.name.startswith(md5_hash):
return f
return None
file_path = await asyncio.to_thread(_find_file)
if file_path is None:
return False
for f in subdir.iterdir():
if f.name.startswith(md5_hash):
await aiofiles.os.remove(f)
return True
return False
await aiofiles.os.remove(file_path)
return True
async def stats(self) -> dict:
"""Get vault statistics."""
if HAS_FILEVAULT and self.filevault:
loop = asyncio.get_event_loop()
return await loop.run_in_executor(
None,
lambda: {
def _get_stats():
return {
'backend': 'filevault',
'count': len(self.filevault) if hasattr(self.filevault, '__len__') else -1
}
)
return await asyncio.to_thread(_get_stats)
else:
count = 0
total_size = 0
for subdir in self.vault_path.iterdir():
if subdir.is_dir():
for f in subdir.iterdir():
count += 1
total_size += f.stat().st_size
def _compute_stats():
count = 0
total_size = 0
for subdir in self.vault_path.iterdir():
if subdir.is_dir():
for f in subdir.iterdir():
count += 1
total_size += f.stat().st_size
return {
'backend': 'directory',
'count': count,
'total_size_bytes': total_size
}
return {
'backend': 'directory',
'count': count,
'total_size_bytes': total_size
}
return await asyncio.to_thread(_compute_stats)

409
tests/unit/test_async_filevault.py Normal file
View file

@ -0,0 +1,409 @@
"""
Unit tests for AsyncFileVault
Tests async wrapper functionality
"""
import pytest
import os
import sys
import tempfile
import shutil
import asyncio
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))
from async_filevault import AsyncVault, create_async_vault
class TestAsyncVaultBasics:
"""Test basic AsyncVault functionality"""
def setup_method(self):
self.temp_dir = tempfile.mkdtemp()
self.vault_path = os.path.join(self.temp_dir, "test_vault")
def teardown_method(self):
if os.path.exists(self.temp_dir):
shutil.rmtree(self.temp_dir)
def test_async_vault_initialization(self):
"""Test AsyncVault initializes correctly"""
vault = AsyncVault(vaultpath=self.vault_path, use_pairs=True)
assert vault.vaultpath == self.vault_path
assert vault.use_pairs is True
assert vault.depth == 3
def test_sync_methods_work(self):
"""Test that sync methods (filename generation) work without await"""
vault = AsyncVault(vaultpath=self.vault_path)
# These should work synchronously
filename = vault.create_filename("test_seed", ".json", absolute=True)
assert filename.startswith(self.vault_path)
assert filename.endswith(".json")
random_filename = vault.create_random_filename(".txt")
assert random_filename.endswith(".txt")
def test_factory_function(self):
"""Test create_async_vault factory function"""
vault = create_async_vault(vaultpath=self.vault_path, salt="test")
assert isinstance(vault, AsyncVault)
class TestAsyncJSONOperations:
"""Test async JSON operations"""
def setup_method(self):
self.temp_dir = tempfile.mkdtemp()
self.vault_path = os.path.join(self.temp_dir, "test_vault")
def teardown_method(self):
if os.path.exists(self.temp_dir):
shutil.rmtree(self.temp_dir)
@pytest.mark.asyncio
async def test_async_atomic_write_json(self):
"""Test async atomic JSON write"""
vault = AsyncVault(vaultpath=self.vault_path)
test_file = vault.create_filename("async_json", ".json", absolute=True)
test_data = {"key": "value", "number": 42}
await vault.atomic_write_json(test_file, test_data)
assert os.path.exists(test_file)
@pytest.mark.asyncio
async def test_async_safe_read_json(self):
"""Test async safe JSON read"""
vault = AsyncVault(vaultpath=self.vault_path)
test_file = vault.create_filename("async_read", ".json", absolute=True)
test_data = {"async": True, "data": [1, 2, 3]}
await vault.atomic_write_json(test_file, test_data)
read_data = await vault.safe_read_json(test_file)
assert read_data == test_data
@pytest.mark.asyncio
async def test_async_safe_read_json_with_default(self):
"""Test async safe JSON read with default"""
vault = AsyncVault(vaultpath=self.vault_path)
test_file = vault.create_filename("nonexistent", ".json", absolute=True)
result = await vault.safe_read_json(test_file, default={"default": True})
assert result == {"default": True}
@pytest.mark.asyncio
async def test_async_write_with_aliases(self):
"""Test async write with aliases"""
vault = AsyncVault(vaultpath=self.vault_path)
test_file = vault.create_filename("primary", ".json", absolute=True)
alias_file = vault.create_filename("alias", ".json", absolute=True)
await vault.atomic_write_json(test_file, {"id": 1}, aliases=["alias"])
# Read through alias
data = await vault.safe_read_json(alias_file)
assert data == {"id": 1}
class TestAsyncTextOperations:
"""Test async text file operations"""
def setup_method(self):
self.temp_dir = tempfile.mkdtemp()
self.vault_path = os.path.join(self.temp_dir, "test_vault")
def teardown_method(self):
if os.path.exists(self.temp_dir):
shutil.rmtree(self.temp_dir)
@pytest.mark.asyncio
async def test_async_write_text_file(self):
"""Test async text file write"""
vault = AsyncVault(vaultpath=self.vault_path)
test_file = vault.create_filename("async_text", ".txt", absolute=True)
test_content = "Async content\nLine 2"
await vault.write_text_file(test_file, test_content)
assert os.path.exists(test_file)
@pytest.mark.asyncio
async def test_async_read_text_file(self):
"""Test async text file read"""
vault = AsyncVault(vaultpath=self.vault_path)
test_file = vault.create_filename("async_read_text", ".txt", absolute=True)
test_content = "Test async read"
await vault.write_text_file(test_file, test_content)
read_content = await vault.read_text_file(test_file)
assert read_content == test_content
@pytest.mark.asyncio
async def test_async_read_text_file_with_default(self):
"""Test async text read with default"""
vault = AsyncVault(vaultpath=self.vault_path)
test_file = vault.create_filename("nonexistent", ".txt", absolute=True)
result = await vault.read_text_file(test_file, default="default content")
assert result == "default content"
@pytest.mark.asyncio
async def test_async_append_text_file(self):
"""Test async text file append"""
vault = AsyncVault(vaultpath=self.vault_path)
test_file = vault.create_filename("async_append", ".txt", absolute=True)
await vault.write_text_file(test_file, "Line 1\n")
await vault.write_text_file(test_file, "Line 2\n", mode="a")
content = await vault.read_text_file(test_file)
assert content == "Line 1\nLine 2\n"
class TestAsyncFileOperations:
"""Test async file existence and removal"""
def setup_method(self):
self.temp_dir = tempfile.mkdtemp()
self.vault_path = os.path.join(self.temp_dir, "test_vault")
def teardown_method(self):
if os.path.exists(self.temp_dir):
shutil.rmtree(self.temp_dir)
@pytest.mark.asyncio
async def test_async_file_exists(self):
"""Test async file existence check"""
vault = AsyncVault(vaultpath=self.vault_path)
test_file = vault.create_filename("exists_test", ".txt", absolute=True)
# File doesn't exist yet
exists_before = await vault.file_exists(test_file)
assert exists_before is False
# Create file
await vault.write_text_file(test_file, "test")
# Now it exists
exists_after = await vault.file_exists(test_file)
assert exists_after is True
@pytest.mark.asyncio
async def test_async_file_exists_with_cache(self):
"""Test async file_exists uses cache when available"""
vault = AsyncVault(vaultpath=self.vault_path, enable_memory_cache=True)
test_file = vault.create_filename("cached", ".txt", absolute=True)
# Manually set cache
vault.mark_file_exists(test_file, True)
# Should return from cache without I/O
exists = await vault.file_exists(test_file)
assert exists is True
@pytest.mark.asyncio
async def test_async_remove_file(self):
"""Test async file removal"""
vault = AsyncVault(vaultpath=self.vault_path)
test_file = vault.create_filename("to_remove", ".txt", absolute=True)
await vault.write_text_file(test_file, "test")
assert os.path.exists(test_file)
result = await vault.remove_file(test_file)
assert result is True
assert not os.path.exists(test_file)
@pytest.mark.asyncio
async def test_async_remove_file_nonexistent(self):
"""Test async remove of non-existent file"""
vault = AsyncVault(vaultpath=self.vault_path)
test_file = vault.create_filename("nonexistent", ".txt", absolute=True)
result = await vault.remove_file(test_file)
assert result is False
@pytest.mark.asyncio
async def test_async_remove_file_if_exists(self):
"""Test async remove_file_if_exists"""
vault = AsyncVault(vaultpath=self.vault_path)
test_file = vault.create_filename("maybe", ".txt", absolute=True)
# Should not raise for non-existent
result1 = await vault.remove_file_if_exists(test_file)
assert result1 is False
# Create and remove
await vault.write_text_file(test_file, "test")
result2 = await vault.remove_file_if_exists(test_file)
assert result2 is True
class TestAsyncConcurrency:
"""Test async concurrent operations"""
def setup_method(self):
self.temp_dir = tempfile.mkdtemp()
self.vault_path = os.path.join(self.temp_dir, "test_vault")
def teardown_method(self):
if os.path.exists(self.temp_dir):
shutil.rmtree(self.temp_dir)
@pytest.mark.asyncio
async def test_concurrent_writes(self):
"""Test multiple concurrent async writes"""
vault = AsyncVault(vaultpath=self.vault_path, use_pairs=True)
async def write_item(i):
filename = vault.create_filename(f"item_{i}", ".json", absolute=True)
await vault.atomic_write_json(filename, {"index": i})
return i
# Run 10 concurrent writes
results = await asyncio.gather(*[write_item(i) for i in range(10)])
assert len(results) == 10
assert set(results) == set(range(10))
# Verify all files exist
for i in range(10):
filename = vault.create_filename(f"item_{i}", ".json", absolute=True)
data = await vault.safe_read_json(filename)
assert data["index"] == i
@pytest.mark.asyncio
async def test_concurrent_reads(self):
"""Test multiple concurrent async reads"""
vault = AsyncVault(vaultpath=self.vault_path)
# Write files first
for i in range(5):
filename = vault.create_filename(f"read_{i}", ".json", absolute=True)
await vault.atomic_write_json(filename, {"value": i * 10})
async def read_item(i):
filename = vault.create_filename(f"read_{i}", ".json", absolute=True)
return await vault.safe_read_json(filename)
# Concurrent reads
results = await asyncio.gather(*[read_item(i) for i in range(5)])
assert len(results) == 5
for i, result in enumerate(results):
assert result["value"] == i * 10
@pytest.mark.asyncio
async def test_mixed_operations(self):
"""Test mixed concurrent operations"""
vault = AsyncVault(vaultpath=self.vault_path, enable_memory_cache=True)
async def operation(i):
filename = vault.create_filename(f"mixed_{i}", ".json", absolute=True)
# Write
await vault.atomic_write_json(filename, {"step": 1, "i": i})
# Read
data = await vault.safe_read_json(filename)
assert data["step"] == 1
# Update
data["step"] = 2
await vault.atomic_write_json(filename, data)
# Verify
final = await vault.safe_read_json(filename)
return final["step"]
results = await asyncio.gather(*[operation(i) for i in range(5)])
assert all(r == 2 for r in results)
class TestAsyncIntegration:
"""Integration tests for async workflows"""
def setup_method(self):
self.temp_dir = tempfile.mkdtemp()
self.vault_path = os.path.join(self.temp_dir, "test_vault")
def teardown_method(self):
if os.path.exists(self.temp_dir):
shutil.rmtree(self.temp_dir)
@pytest.mark.asyncio
async def test_full_async_workflow(self):
"""Test complete async CRUD workflow"""
vault = AsyncVault(
vaultpath=self.vault_path,
depth=5,
use_pairs=True,
enable_memory_cache=True
)
key = "workflow_key"
filename = vault.create_filename(key, ".json", absolute=True)
# Create
await vault.atomic_write_json(filename, {"version": 1})
assert await vault.file_exists(filename)
# Read
data = await vault.safe_read_json(filename)
assert data["version"] == 1
# Update
data["version"] = 2
data["updated"] = True
await vault.atomic_write_json(filename, data)
# Verify
updated = await vault.safe_read_json(filename)
assert updated["version"] == 2
assert updated["updated"] is True
# Delete
await vault.remove_file(filename)
assert not await vault.file_exists(filename)
@pytest.mark.asyncio
async def test_async_pairs_mode(self):
"""Test async with pairs mode (neopig use case)"""
vault = AsyncVault(
vaultpath=self.vault_path,
depth=9,
use_pairs=True,
salt="test_salt"
)
# Simulate domain-based storage
domains = ["example.com", "test.org", "sample.net"]
for domain in domains:
filename = vault.create_filename(domain, ".json", absolute=True)
await vault.atomic_write_json(filename, {"domain": domain, "crawled": True})
# Verify all
for domain in domains:
filename = vault.create_filename(domain, ".json", absolute=True)
data = await vault.safe_read_json(filename)
assert data["domain"] == domain
assert data["crawled"] is True

1031
tests/unit/test_filevault.py Normal file
View file

File diff suppressed because it is too large Load diff