esbi-cli 0.2.1__py3-none-any.whl

esbi_cli/netguard.py

@@ -0,0 +1,127 @@
+"""Refuse to fetch URLs that could reach internal services (SSRF): email links are untrusted."""
+
+import ipaddress
+import socket
+import time
+from collections.abc import Callable
+from urllib.parse import urljoin, urlparse
+
+import httpx
+
+
+class UnsafeURL(ValueError):
+    pass
+
+
+class CannotResolve(UnsafeURL):
+    """The host does not resolve at all (a dead domain): unreachable, not "refused"."""
+
+
+NAT64 = ipaddress.ip_network("64:ff9b::/96")
+
+
+def _is_public(address: str) -> bool:
+    ip = ipaddress.ip_address(address.split("%")[0])
+    if isinstance(ip, ipaddress.IPv6Address) and ip.ipv4_mapped:
+        ip = ip.ipv4_mapped
+    if isinstance(ip, ipaddress.IPv6Address) and ip in NAT64:  # judge the IPv4 address it wraps
+        ip = ipaddress.IPv4Address(int(ip) & 0xFFFFFFFF)
+    return ip.is_global and not ip.is_multicast
+
+
+def check_public_url(url: str, resolver: Callable = socket.getaddrinfo) -> list[str]:
+    """Raise UnsafeURL unless url is http(s) and its host resolves only to public addresses.
+    Returns those addresses: connect to one of them, never to the name (see safe_get)."""
+    parsed = urlparse(url)
+    if parsed.scheme not in ("http", "https"):
+        raise UnsafeURL(f"Only http(s) URLs are fetched, not {parsed.scheme or 'this'!r}")
+    host = parsed.hostname
+    if not host:
+        raise UnsafeURL("The URL has no host")
+    try:
+        addresses = [str(ipaddress.ip_address(host))]  # an IP literal needs no DNS
+    except ValueError:
+        try:
+            infos = resolver(host, parsed.port or 443, type=socket.SOCK_STREAM)
+        except OSError as exc:
+            raise CannotResolve(f"Cannot resolve {host}: {exc}") from None
+        addresses = list(dict.fromkeys(info[4][0] for info in infos))
+    for address in addresses:
+        if not _is_public(address):
+            raise UnsafeURL(f"{host} resolves to a non-public address ({address})")
+    return addresses
+
+
+# Behind a proxy it is the proxy that resolves and connects, so the address check below would no
+# longer describe where the request goes. Environment proxies (and the other environment settings
+# httpx reads) are therefore ignored unless the user opts in with [network] use_environment_proxy.
+use_environment_proxy = False
+
+MAX_REDIRECTS = 5
+MAX_BYTES = 20_000_000  # a page or PDF bigger than this is not a note
+DEADLINE_SECONDS = 120  # one whole download: a server that trickles bytes cannot hold a run
+
+
+def _read_capped(response: httpx.Response, url: str) -> httpx.Response:
+    body, started = bytearray(), time.monotonic()
+    for chunk in response.iter_bytes():
+        body += chunk
+        if len(body) > MAX_BYTES:
+            raise UnsafeURL(f"The response is larger than {MAX_BYTES // 1_000_000} MB")
+        if time.monotonic() - started > DEADLINE_SECONDS:
+            raise UnsafeURL(f"The download took more than {DEADLINE_SECONDS} s")
+    return httpx.Response(
+        response.status_code,
+        headers={k: v for k, v in response.headers.items() if k.lower() != "content-encoding"},
+        content=bytes(body),
+        request=httpx.Request("GET", url),  # the URL the caller asked for, not the pinned IP
+    )
+
+
+def _fetch_pinned(client: httpx.Client, url: str, addresses: list[str]) -> httpx.Response:
+    target = httpx.URL(url)
+    for i, address in enumerate(addresses):
+        try:
+            with client.stream(
+                "GET",
+                target.copy_with(host=address),
+                headers={"Host": target.netloc.decode()},
+                extensions={"sni_hostname": target.host},
+                follow_redirects=False,
+            ) as response:
+                if response.is_redirect:
+                    return httpx.Response(response.status_code, headers=response.headers)
+                return _read_capped(response, url)
+        except httpx.ConnectError:
+            if i == len(addresses) - 1:  # tried them all
+                raise
+
+
+def safe_get(
+    url: str,
+    *,
+    client: httpx.Client | None = None,
+    resolver: Callable = socket.getaddrinfo,
+    headers: dict | None = None,
+    timeout_seconds: float = 30,
+) -> httpx.Response:
+    """GET a URL, following redirects manually so every hop passes check_public_url.
+
+    The name is resolved once per hop, and the request goes to one of the addresses that passed
+    the check (the URL carries the IP; Host and the TLS server name stay the original host, so
+    certificates are still verified for it). A second lookup by the HTTP client, which a hostile
+    DNS server could answer with a private address, never happens."""
+    own_client = client is None
+    client = client or httpx.Client(
+        headers=headers, timeout=timeout_seconds, trust_env=use_environment_proxy
+    )
+    try:
+        for _ in range(MAX_REDIRECTS + 1):
+            response = _fetch_pinned(client, url, check_public_url(url, resolver=resolver))
+            if not response.is_redirect:
+                return response
+            url = urljoin(url, response.headers["location"])
+        raise UnsafeURL(f"Too many redirects (more than {MAX_REDIRECTS})")
+    finally:
+        if own_client:
+            client.close()

esbi_cli/privacy.py

@@ -0,0 +1,81 @@
+"""Email stays on the machine: it is never sent to a model that sends text away."""
+
+import re
+from urllib.parse import urlparse
+
+from esbi_cli import lang
+from esbi_cli.llm.adapter import is_loopback
+from esbi_cli.vault import Vault
+
+
+class PrivacyError(ValueError):
+    pass
+
+
+def sends_text_out(llm) -> bool:
+    """True for models run by someone else (API, subscription); a local model, and any fake, is not."""
+    return bool(getattr(llm, "sends_text_out", False))
+
+
+def remote_host(model: str, base_url: str | None) -> str | None:
+    """The host serving an Ollama or LM Studio model when that is not this machine, else None."""
+    if model.partition("/")[0] not in ("ollama", "lmstudio") or not base_url:
+        return None
+    return None if is_loopback(base_url) else urlparse(base_url).hostname
+
+
+def remote_warning(model: str, host: str) -> str:
+    return (
+        f"{model} is served by {host}, not this machine: the text of your notes goes to that "
+        "host, and email is kept off it."
+    )
+
+
+def section_pattern(source_title: str) -> re.Pattern:
+    """The `## From [[source]]` section a source adds to a concept page, in any catalogued language
+    (a vault may hold sections written before the language setting changed)."""
+    link = re.escape(f"[[{source_title}]]")
+    heads = "|".join(re.escape(h) for h in lang.every("from_source"))
+    return re.compile(rf"\n*## (?:{heads}) {link}\n.*?(?=\n## |\Z)", re.DOTALL)
+
+
+def private_sources(vault: Vault) -> set[str]:
+    return {p.title for p in vault.iter_pages(("sources",)) if p.meta.get("kind") == "email"}
+
+
+def _email_links(vault: Vault) -> set[str]:
+    return {f"[[{t}]]" for t in private_sources(vault)}
+
+
+def private_titles(vault: Vault) -> set[str]:
+    """Pages to hide from a cloud model: email sources, pages built only from them, and every
+    synthesis that used any email page (its text is an answer written from that page)."""
+    emails = private_sources(vault)
+    linked = _email_links(vault)
+    hidden = set(emails)
+    for page in vault.iter_pages(("concepts", "entities", "syntheses")):
+        sources = page.meta.get("sources") or []
+        if page.kind == "syntheses":
+            if any(s in linked for s in sources):
+                hidden.add(page.title)
+        elif sources and all(s in linked for s in sources):
+            hidden.add(page.title)
+    return hidden
+
+
+def email_touched(vault: Vault) -> set[str]:
+    """Pages email helped write (some source of theirs is an email): their one-line summaries may
+    carry email text, so a cloud model is offered the title only."""
+    linked = _email_links(vault)
+    return {
+        p.title
+        for p in vault.iter_pages(("concepts", "entities", "syntheses"))
+        if any(s in linked for s in (p.meta.get("sources") or []))
+    }
+
+
+def public_body(body: str, vault_private_sources: set[str]) -> str:
+    """A page's text without the sections that came from email."""
+    for title in vault_private_sources:
+        body = section_pattern(title).sub("", body)
+    return body.strip()

esbi_cli/queue.py

@@ -0,0 +1,179 @@
+"""Persistent work queue of sources waiting to be ingested (SQLite)."""
+
+import sqlite3
+from collections.abc import Collection
+from dataclasses import dataclass
+from datetime import date
+from pathlib import Path
+from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse
+
+_TRACKING_PARAMS = frozenset({"fbclid", "gclid", "mc_cid", "mc_eid", "trk", "trkemail"})
+
+
+def normalize_target(target: str) -> str:
+    """Canonical form used to detect duplicates: URLs lose tracking params and fragments."""
+    parsed = urlparse(target)
+    if parsed.scheme not in ("http", "https"):
+        return str(Path(target).expanduser().resolve())
+    query = [
+        (k, v)
+        for k, v in parse_qsl(parsed.query, keep_blank_values=True)
+        if not k.lower().startswith("utm_") and k.lower() not in _TRACKING_PARAMS
+    ]
+    path = parsed.path.rstrip("/") or "/"
+    return urlunparse(
+        (parsed.scheme, parsed.netloc.lower(), path, "", urlencode(sorted(query)), "")
+    )
+
+
+_COLUMNS = "id, target, origin, attempts, error, captured, label"
+
+
+def _item(row: tuple) -> "Item":
+    id_, target, origin, attempts, error, captured, label = row
+    return Item(
+        id_,
+        target,
+        origin,
+        attempts,
+        error,
+        date.fromisoformat(captured) if captured else None,
+        label,
+    )
+
+
+@dataclass
+class Item:
+    id: int
+    target: str
+    origin: str
+    attempts: int = 0
+    error: str | None = None
+    captured: date | None = None  # when the source was saved, if known (legacy import)
+    label: str | None = None  # human-readable name (clip title, file name)
+
+
+class Queue:
+    def __init__(self, path: Path, max_attempts: int = 3):
+        self.max_attempts = max_attempts
+        path.parent.mkdir(parents=True, exist_ok=True)
+        self._db = sqlite3.connect(path)
+        self._db.execute(
+            """CREATE TABLE IF NOT EXISTS items (
+                id INTEGER PRIMARY KEY,
+                target TEXT NOT NULL UNIQUE,
+                origin TEXT NOT NULL,
+                status TEXT NOT NULL DEFAULT 'queued',
+                attempts INTEGER NOT NULL DEFAULT 0,
+                error TEXT,
+                added_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
+                captured TEXT,
+                label TEXT
+            )"""
+        )
+        # databases created before these columns existed keep working
+        existing = {row[1] for row in self._db.execute("PRAGMA table_info(items)")}
+        for column in ("captured", "label"):
+            if column not in existing:
+                self._db.execute(f"ALTER TABLE items ADD COLUMN {column} TEXT")
+        self._db.commit()
+
+    def add(
+        self,
+        target: str,
+        origin: str,
+        captured: date | None = None,
+        label: str | None = None,
+    ) -> bool:
+        """Queue a source. Returns False if an equivalent one is already known."""
+        cur = self._db.execute(
+            "INSERT OR IGNORE INTO items (target, origin, captured, label) VALUES (?, ?, ?, ?)",
+            (normalize_target(target), origin, captured.isoformat() if captured else None, label),
+        )
+        self._db.commit()
+        return cur.rowcount == 1
+
+    def counts(self) -> dict[str, int]:
+        rows = self._db.execute("SELECT status, COUNT(*) FROM items GROUP BY status").fetchall()
+        return dict(rows)
+
+    def claim(self, n: int, exclude: Collection[int] = ()) -> list[Item]:
+        """Take up to n of the oldest queued items (skipping `exclude`), marking them processing."""
+        skip = ",".join("?" for _ in exclude)
+        rows = self._db.execute(
+            f"SELECT {_COLUMNS} FROM items WHERE status = 'queued'"
+            + (f" AND id NOT IN ({skip})" if skip else "")
+            + " ORDER BY id LIMIT ?",
+            (*exclude, n),
+        ).fetchall()
+        self._db.executemany(
+            "UPDATE items SET status = 'processing' WHERE id = ?", [(r[0],) for r in rows]
+        )
+        self._db.commit()
+        return [_item(r) for r in rows]
+
+    def complete(self, item_id: int) -> None:
+        self._db.execute("UPDATE items SET status = 'done' WHERE id = ?", (item_id,))
+        self._db.commit()
+
+    def fail(self, item_id: int, error: str) -> None:
+        """Record a failed attempt: retry next run, or park as 'failed' after max_attempts."""
+        self._db.execute(
+            """UPDATE items SET attempts = attempts + 1, error = ?,
+               status = CASE WHEN attempts + 1 >= ? THEN 'failed' ELSE 'queued' END
+               WHERE id = ?""",
+            (error, self.max_attempts, item_id),
+        )
+        self._db.commit()
+
+    def items(self, status: str) -> list[Item]:
+        rows = self._db.execute(
+            f"SELECT {_COLUMNS} FROM items WHERE status = ? ORDER BY id", (status,)
+        ).fetchall()
+        return [_item(r) for r in rows]
+
+    def get(self, target: str) -> Item | None:
+        row = self._db.execute(
+            f"SELECT {_COLUMNS} FROM items WHERE target = ?", (normalize_target(target),)
+        ).fetchone()
+        return _item(row) if row else None
+
+    def requeue_failed(self, target: str | None = None) -> int:
+        """Put parked failures back in the queue with a fresh attempt count. Returns how many."""
+        sql = (
+            "UPDATE items SET status = 'queued', attempts = 0, error = NULL WHERE status = 'failed'"
+        )
+        args: tuple = ()
+        if target is not None:
+            sql, args = sql + " AND target = ?", (normalize_target(target),)
+        cur = self._db.execute(sql, args)
+        self._db.commit()
+        return cur.rowcount
+
+    def remove(self, target: str) -> bool:
+        """Delete an item whatever its state. The original file, if any, is not touched."""
+        cur = self._db.execute("DELETE FROM items WHERE target = ?", (normalize_target(target),))
+        self._db.commit()
+        return cur.rowcount == 1
+
+    def recover(self) -> int:
+        """Requeue items a previous run claimed but never finished. Returns how many. The lost run
+        counts as an attempt: an item that kills the process (a PDF that exhausts memory) is parked
+        after max_attempts instead of killing every run."""
+        cur = self._db.execute(
+            """UPDATE items SET attempts = attempts + 1,
+               status = CASE WHEN attempts + 1 >= ? THEN 'failed' ELSE 'queued' END,
+               error = CASE WHEN attempts + 1 >= ? THEN 'the run died while reading this' ELSE error END
+               WHERE status = 'processing'""",
+            (self.max_attempts, self.max_attempts),
+        )
+        self._db.commit()
+        return cur.rowcount
+
+    def release(self, item_id: int) -> bool:
+        """Give back a claimed item untouched (no attempt counted). False if it was not claimed."""
+        cur = self._db.execute(
+            "UPDATE items SET status = 'queued' WHERE id = ? AND status = 'processing'", (item_id,)
+        )
+        self._db.commit()
+        return cur.rowcount == 1

esbi_cli/reingest.py

@@ -0,0 +1,165 @@
+"""`sb reingest`: rebuild source notes from the saved raw files, keeping what the user did to them."""
+
+import shutil
+import subprocess
+import time
+from collections.abc import Callable
+from contextlib import ExitStack
+from dataclasses import dataclass, field
+from datetime import date
+
+from esbi_cli.config import Config
+from esbi_cli.extract import ExtractedDoc, ExtractError
+from esbi_cli.extract.image import IMAGE_SUFFIXES, image_figure, to_png
+from esbi_cli.extract.pdf import extract_pdf_bytes
+from esbi_cli.gitops import commit_vault
+from esbi_cli.ingest.apply import NOTE_FORMAT
+from esbi_cli.ingest.pipeline import ingest
+from esbi_cli.interrupts import deferred
+from esbi_cli.llm.adapter import LLM, LLMError, LLMTimeout
+from esbi_cli.privacy import section_pattern
+from esbi_cli.vault import Page, Vault, fold, slugify
+
+KEPT = ("status", "read", "content_hash")  # what the user (or dedupe) owns: never rebuilt
+
+
+@dataclass
+class ReingestResult:
+    done: list[str] = field(default_factory=list)
+    skipped: list[tuple[str, str]] = field(default_factory=list)  # (title, why)
+    failed: list[tuple[str, str]] = field(default_factory=list)
+    warnings: list[tuple[str, str]] = field(default_factory=list)  # (title, what went wrong)
+    stopped: bool = False  # the model went away; run again to resume
+    tag: str | None = None
+
+
+def _raw_doc(vault: Vault, note: Page) -> ExtractedDoc:
+    """The source as it was saved: the PDF (re-extracted, so figures come back) or the snapshot."""
+    raw = vault.root / str(note.meta.get("raw") or "")
+    if not note.meta.get("raw") or not raw.is_file():
+        raise ExtractError(f"the file is missing from raw/ ({note.meta.get('raw')})")
+    snap = vault.read_page(raw)
+    doc = ExtractedDoc(
+        note.title, snap.body, note.meta.get("kind") or "article", note.meta.get("url")
+    )
+    pdf = raw.with_suffix(".pdf")
+    if pdf.is_file():
+        data = pdf.read_bytes()
+        try:
+            doc = extract_pdf_bytes(data, note.title)
+            doc.url = note.meta.get("url")
+        except ExtractError:  # a scanned PDF: the snapshot holds the text the OCR model read
+            doc.pdf_bytes = data  # still tells the privacy rules whether it came from a mail
+    for suffix in IMAGE_SUFFIXES:  # a picture note: show the original again, no new OCR
+        if (image := raw.with_suffix(suffix)).is_file():
+            doc.figures = [image_figure(to_png(image.read_bytes(), image.name))]
+    return doc
+
+
+def _forget(vault: Vault, note: Page) -> None:
+    """Remove what this source added elsewhere; concept pages nobody else cites are deleted."""
+    link = f"[[{note.title}]]"
+    section = section_pattern(note.title)
+    for page in vault.iter_pages(("concepts", "entities", "syntheses")):
+        sources = list(page.meta.get("sources") or [])
+        if link not in sources:
+            continue
+        sources.remove(link)
+        if not sources and page.meta.get("type") in ("concept", "entity"):
+            page.path.unlink()
+            continue
+        page.meta["sources"] = sources
+        page.body = section.sub("", page.body).strip()
+        vault.write_page(page)
+    note.path.unlink()
+    shutil.rmtree(vault.root / "attachments" / slugify(note.title), ignore_errors=True)
+
+
+def _tag_before(vault: Vault, today: date) -> str:
+    """One command to undo the whole rebuild: `git reset --hard <tag>` in the vault."""
+    commit_vault(vault.root, "chore: state before reingest")
+    tag = f"pre-reingest-{today.isoformat()}"
+    # an existing tag means this is a resumed run: the first tag is the true "before"
+    subprocess.run(["git", "tag", tag], cwd=vault.root, capture_output=True)
+    return tag
+
+
+def reingest_all(
+    vault: Vault,
+    llm: LLM,
+    synth_llm: LLM | None,
+    cfg: Config,
+    *,
+    today: date | None = None,
+    all_sources: bool = False,
+    private_llm: LLM | None = None,
+    only: list[str] | None = None,
+    on_progress: Callable[[str], None] = lambda _: None,
+) -> ReingestResult:
+    """Rebuild every note that is not yet in the current format (or all of them), one commit each.
+
+    Notes carry `format: 2` once rebuilt, so an interrupted run resumes where it stopped.
+    """
+    today = today or date.today()
+    result = ReingestResult()
+    wanted = [fold(t) for t in only or []]  # naming notes means "these, whatever their format"
+
+    def selected(note: Page) -> bool:
+        if wanted:
+            return any(w in fold(note.title) for w in wanted)
+        return all_sources or note.meta.get("format") != NOTE_FORMAT
+
+    todo = [n for n in vault.iter_pages(("sources",)) if selected(n)]
+    if not todo:
+        return result
+    result.tag = _tag_before(vault, today)
+    started = time.monotonic()
+    for n, note in enumerate(todo, 1):
+        left = ""
+        if n > 1:  # the pace of the notes done so far is the best guess for the rest
+            minutes = (time.monotonic() - started) / (n - 1) * (len(todo) - n + 1) / 60
+            left = f" (~{minutes:.0f} min left)"
+        on_progress(f"[{n}/{len(todo)}] {note.title}{left}")
+        try:
+            doc = _raw_doc(vault, note)
+        except ExtractError as exc:
+            result.skipped.append((note.title, str(exc)))
+            continue
+        with ExitStack() as writing:
+            try:
+                done = ingest(
+                    note.title,
+                    vault=vault,
+                    llm=llm,
+                    synth_llm=synth_llm,
+                    private_llm=private_llm,
+                    cfg=cfg,
+                    force=True,
+                    extractor=lambda _, d=doc: d,
+                    today=today,
+                    on_step=lambda step: on_progress(f"    ... {step}"),
+                    keep_title=note.title,
+                    raw_path=vault.root / note.meta["raw"],
+                    # from the moment the old note goes until the new one is restored and committed,
+                    # a signal waits: stopping in between would lose the user's read state
+                    before_write=lambda note=note: (
+                        writing.enter_context(deferred()),
+                        _forget(vault, note),
+                    ),
+                )
+            except Exception as exc:  # one bad source must not stop the rebuild
+                if isinstance(exc, LLMError) and not isinstance(exc, LLMTimeout):
+                    result.stopped = True  # the model is the problem, not this source: resume later
+                    break
+                result.failed.append((note.title, f"{type(exc).__name__}: {exc}"))
+                continue
+            result.warnings += [(note.title, w) for w in done.warnings]
+            rebuilt = vault.read_page(done.applied.source_path)
+            for key in KEPT:
+                rebuilt.meta[key] = note.meta.get(key)
+            rebuilt.meta["captured"] = note.meta.get("captured", rebuilt.meta["captured"])
+            rebuilt.meta["tags"] = sorted({*(note.meta.get("tags") or []), *rebuilt.meta["tags"]})
+            vault.write_page(rebuilt)
+            commit_vault(vault.root, f"reingest: {done.applied.source_title}")
+            result.done.append(done.applied.source_title)
+    return result