Zakaria
418 lines
13 KiB
Python
418 lines
13 KiB
Python
"""Per-run notes storage — mirrored to {state_dir}/notes.json."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import tempfile
|
|
import threading
|
|
import uuid
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from agents import RunContextWrapper, function_tool
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
_notes_storage: dict[str, dict[str, Any]] = {}
|
|
_VALID_NOTE_CATEGORIES = ["general", "findings", "methodology", "questions", "plan", "wiki"]
|
|
_notes_lock = threading.RLock()
|
|
_DEFAULT_CONTENT_PREVIEW_CHARS = 280
|
|
|
|
_notes_path: Path | None = None
|
|
|
|
|
|
def hydrate_notes_from_disk(state_dir: Path) -> None:
|
|
global _notes_path # noqa: PLW0603
|
|
_notes_path = state_dir / "notes.json"
|
|
with _notes_lock:
|
|
_notes_storage.clear()
|
|
if not _notes_path.exists():
|
|
return
|
|
try:
|
|
data = json.loads(_notes_path.read_text(encoding="utf-8"))
|
|
except (OSError, json.JSONDecodeError):
|
|
logger.exception(
|
|
"notes.json at %s is unreadable; starting with empty notes",
|
|
_notes_path,
|
|
)
|
|
return
|
|
if not isinstance(data, dict):
|
|
return
|
|
_notes_storage.update(
|
|
{
|
|
nid: note
|
|
for nid, note in data.items()
|
|
if isinstance(nid, str) and isinstance(note, dict)
|
|
}
|
|
)
|
|
logger.info(
|
|
"notes hydrated from %s (%d note(s))",
|
|
_notes_path,
|
|
len(_notes_storage),
|
|
)
|
|
|
|
|
|
def _persist() -> None:
|
|
path = _notes_path
|
|
if path is None:
|
|
return
|
|
try:
|
|
payload = json.dumps(_notes_storage, ensure_ascii=False, default=str)
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
with (
|
|
_notes_lock,
|
|
tempfile.NamedTemporaryFile(
|
|
mode="w",
|
|
encoding="utf-8",
|
|
dir=str(path.parent),
|
|
prefix=f".{path.name}.",
|
|
suffix=".tmp",
|
|
delete=False,
|
|
) as tmp,
|
|
):
|
|
tmp.write(payload)
|
|
tmp_path = Path(tmp.name)
|
|
tmp_path.replace(path)
|
|
except Exception:
|
|
logger.exception("notes persist to %s failed", path)
|
|
|
|
|
|
def _filter_notes(
|
|
category: str | None = None,
|
|
tags: list[str] | None = None,
|
|
search_query: str | None = None,
|
|
) -> list[dict[str, Any]]:
|
|
filtered: list[dict[str, Any]] = []
|
|
for note_id, note in _notes_storage.items():
|
|
if category and note.get("category") != category:
|
|
continue
|
|
if tags:
|
|
note_tags = note.get("tags", [])
|
|
if not any(tag in note_tags for tag in tags):
|
|
continue
|
|
if search_query:
|
|
search_lower = search_query.lower()
|
|
title_match = search_lower in note.get("title", "").lower()
|
|
content_match = search_lower in note.get("content", "").lower()
|
|
if not (title_match or content_match):
|
|
continue
|
|
entry = note.copy()
|
|
entry["note_id"] = note_id
|
|
filtered.append(entry)
|
|
filtered.sort(key=lambda x: x.get("created_at", ""), reverse=True)
|
|
return filtered
|
|
|
|
|
|
def _to_note_listing_entry(
|
|
note: dict[str, Any],
|
|
*,
|
|
include_content: bool = False,
|
|
) -> dict[str, Any]:
|
|
entry = {
|
|
"note_id": note.get("note_id"),
|
|
"title": note.get("title", ""),
|
|
"category": note.get("category", "general"),
|
|
"tags": note.get("tags", []),
|
|
"created_at": note.get("created_at", ""),
|
|
"updated_at": note.get("updated_at", ""),
|
|
}
|
|
content = str(note.get("content", ""))
|
|
if include_content:
|
|
entry["content"] = content
|
|
elif content:
|
|
if len(content) > _DEFAULT_CONTENT_PREVIEW_CHARS:
|
|
entry["content_preview"] = f"{content[:_DEFAULT_CONTENT_PREVIEW_CHARS].rstrip()}..."
|
|
else:
|
|
entry["content_preview"] = content
|
|
return entry
|
|
|
|
|
|
def _create_note_impl(
|
|
title: str,
|
|
content: str,
|
|
category: str = "general",
|
|
tags: list[str] | None = None,
|
|
) -> dict[str, Any]:
|
|
with _notes_lock:
|
|
try:
|
|
if not title or not title.strip():
|
|
return {"success": False, "error": "Title cannot be empty", "note_id": None}
|
|
if not content or not content.strip():
|
|
return {"success": False, "error": "Content cannot be empty", "note_id": None}
|
|
if category not in _VALID_NOTE_CATEGORIES:
|
|
return {
|
|
"success": False,
|
|
"error": (
|
|
f"Invalid category. Must be one of: {', '.join(_VALID_NOTE_CATEGORIES)}"
|
|
),
|
|
"note_id": None,
|
|
}
|
|
|
|
note_id = str(uuid.uuid4())[:6]
|
|
|
|
timestamp = datetime.now(UTC).isoformat()
|
|
note = {
|
|
"title": title.strip(),
|
|
"content": content.strip(),
|
|
"category": category,
|
|
"tags": tags or [],
|
|
"created_at": timestamp,
|
|
"updated_at": timestamp,
|
|
}
|
|
_notes_storage[note_id] = note
|
|
except (ValueError, TypeError) as e:
|
|
return {"success": False, "error": f"Failed to create note: {e}", "note_id": None}
|
|
else:
|
|
_persist()
|
|
return {
|
|
"success": True,
|
|
"note_id": note_id,
|
|
"message": f"Note '{title}' created successfully",
|
|
"total_count": len(_notes_storage),
|
|
}
|
|
|
|
|
|
def _list_notes_impl(
|
|
category: str | None = None,
|
|
tags: list[str] | None = None,
|
|
search: str | None = None,
|
|
include_content: bool = False,
|
|
) -> dict[str, Any]:
|
|
with _notes_lock:
|
|
try:
|
|
filtered = _filter_notes(category=category, tags=tags, search_query=search)
|
|
notes = [_to_note_listing_entry(n, include_content=include_content) for n in filtered]
|
|
except (ValueError, TypeError) as e:
|
|
return {
|
|
"success": False,
|
|
"error": f"Failed to list notes: {e}",
|
|
"notes": [],
|
|
"filtered_count": 0,
|
|
"total_count": 0,
|
|
}
|
|
return {
|
|
"success": True,
|
|
"notes": notes,
|
|
"filtered_count": len(notes),
|
|
"total_count": len(_notes_storage),
|
|
}
|
|
|
|
|
|
def _get_note_impl(note_id: str) -> dict[str, Any]:
|
|
with _notes_lock:
|
|
try:
|
|
if not note_id or not note_id.strip():
|
|
return {"success": False, "error": "Note ID cannot be empty", "note": None}
|
|
note = _notes_storage.get(note_id)
|
|
if note is None:
|
|
return {
|
|
"success": False,
|
|
"error": f"Note with ID '{note_id}' not found",
|
|
"note": None,
|
|
}
|
|
note_with_id = note.copy()
|
|
note_with_id["note_id"] = note_id
|
|
except (ValueError, TypeError) as e:
|
|
return {"success": False, "error": f"Failed to get note: {e}", "note": None}
|
|
else:
|
|
return {"success": True, "note": note_with_id}
|
|
|
|
|
|
def _update_note_impl(
|
|
note_id: str,
|
|
title: str | None = None,
|
|
content: str | None = None,
|
|
tags: list[str] | None = None,
|
|
) -> dict[str, Any]:
|
|
with _notes_lock:
|
|
try:
|
|
if note_id not in _notes_storage:
|
|
return {"success": False, "error": f"Note with ID '{note_id}' not found"}
|
|
note = _notes_storage[note_id]
|
|
if title is not None:
|
|
if not title.strip():
|
|
return {"success": False, "error": "Title cannot be empty"}
|
|
note["title"] = title.strip()
|
|
if content is not None:
|
|
if not content.strip():
|
|
return {"success": False, "error": "Content cannot be empty"}
|
|
note["content"] = content.strip()
|
|
if tags is not None:
|
|
note["tags"] = tags
|
|
note["updated_at"] = datetime.now(UTC).isoformat()
|
|
except (ValueError, TypeError) as e:
|
|
return {"success": False, "error": f"Failed to update note: {e}"}
|
|
else:
|
|
_persist()
|
|
return {
|
|
"success": True,
|
|
"note_id": note_id,
|
|
"message": f"Note '{note['title']}' updated successfully",
|
|
"total_count": len(_notes_storage),
|
|
}
|
|
|
|
|
|
def _delete_note_impl(note_id: str) -> dict[str, Any]:
|
|
with _notes_lock:
|
|
try:
|
|
if note_id not in _notes_storage:
|
|
return {"success": False, "error": f"Note with ID '{note_id}' not found"}
|
|
note = _notes_storage[note_id]
|
|
note_title = note["title"]
|
|
del _notes_storage[note_id]
|
|
except (ValueError, TypeError) as e:
|
|
return {"success": False, "error": f"Failed to delete note: {e}"}
|
|
else:
|
|
_persist()
|
|
return {
|
|
"success": True,
|
|
"note_id": note_id,
|
|
"message": f"Note '{note_title}' deleted successfully",
|
|
"total_count": len(_notes_storage),
|
|
}
|
|
|
|
|
|
@function_tool(timeout=30)
|
|
async def create_note(
|
|
ctx: RunContextWrapper,
|
|
title: str,
|
|
content: str,
|
|
category: str = "general",
|
|
tags: list[str] | None = None,
|
|
) -> str:
|
|
"""Document an observation, finding, methodology step, or research note.
|
|
|
|
Notes are visible to every agent in the same scan for the lifetime
|
|
of the run; they live in-memory only and are cleared when the
|
|
process exits.
|
|
|
|
For actionable tasks, use ``todo`` instead — notes are for capturing
|
|
information, todos are for tracking work.
|
|
|
|
Categories:
|
|
|
|
- ``general`` — default, anything that doesn't fit elsewhere.
|
|
- ``findings`` — confirmed vulnerabilities or weaknesses (write
|
|
these up promptly; you'll cite them when filing reports).
|
|
- ``methodology`` — what you tried, what worked, what didn't —
|
|
useful for the final scan report.
|
|
- ``questions`` — open questions / things to come back to.
|
|
- ``plan`` — multi-step plans you want to track.
|
|
- ``wiki`` — long-form repository or target maps.
|
|
|
|
Tags are free-form (e.g. ``["sqli", "auth", "critical"]``) — useful
|
|
for later ``list_notes(tags=...)`` filtering.
|
|
|
|
Args:
|
|
title: Short headline.
|
|
content: Full note body. Markdown is preserved.
|
|
category: One of the categories above. Default ``"general"``.
|
|
tags: Optional free-form tags.
|
|
"""
|
|
return json.dumps(
|
|
await asyncio.to_thread(_create_note_impl, title, content, category, tags),
|
|
ensure_ascii=False,
|
|
default=str,
|
|
)
|
|
|
|
|
|
@function_tool(timeout=30)
|
|
async def list_notes(
|
|
ctx: RunContextWrapper,
|
|
category: str | None = None,
|
|
tags: list[str] | None = None,
|
|
search: str | None = None,
|
|
include_content: bool = False,
|
|
) -> str:
|
|
"""List existing notes — metadata-first by default.
|
|
|
|
Filters compose: passing ``category="findings"`` and
|
|
``tags=["sqli"]`` returns notes that are *both* in the findings
|
|
category AND have at least one of those tags.
|
|
|
|
By default each entry includes a ``content_preview`` (first 280
|
|
chars). Set ``include_content=True`` to get full bodies — useful
|
|
when you need to scan many notes; expensive in tokens for large
|
|
notes.
|
|
|
|
Args:
|
|
category: Filter by category.
|
|
tags: Filter to notes that have any of these tags.
|
|
search: Substring match against title and content.
|
|
include_content: When False (default) entries have a preview;
|
|
when True the full ``content`` is included.
|
|
"""
|
|
return json.dumps(
|
|
await asyncio.to_thread(
|
|
_list_notes_impl,
|
|
category=category,
|
|
tags=tags,
|
|
search=search,
|
|
include_content=include_content,
|
|
),
|
|
ensure_ascii=False,
|
|
default=str,
|
|
)
|
|
|
|
|
|
@function_tool(timeout=30)
|
|
async def get_note(ctx: RunContextWrapper, note_id: str) -> str:
|
|
"""Fetch one note by its 6-char ID. Returns the full content.
|
|
|
|
Args:
|
|
note_id: Note id from ``create_note`` or a ``list_notes`` entry.
|
|
"""
|
|
return json.dumps(
|
|
await asyncio.to_thread(_get_note_impl, note_id), ensure_ascii=False, default=str
|
|
)
|
|
|
|
|
|
@function_tool(timeout=30)
|
|
async def update_note(
|
|
ctx: RunContextWrapper,
|
|
note_id: str,
|
|
title: str | None = None,
|
|
content: str | None = None,
|
|
tags: list[str] | None = None,
|
|
) -> str:
|
|
"""Update a note's title, content, or tags.
|
|
|
|
Pass ``None`` for any field you want left unchanged. Replacing
|
|
``content`` is a full overwrite — to append, fetch first with
|
|
``get_note``, concat, and pass the result.
|
|
|
|
Args:
|
|
note_id: Target note's 6-char ID.
|
|
title: New title, or ``None`` to keep.
|
|
content: New content, or ``None`` to keep.
|
|
tags: New tags list, or ``None`` to keep.
|
|
"""
|
|
return json.dumps(
|
|
await asyncio.to_thread(
|
|
_update_note_impl,
|
|
note_id=note_id,
|
|
title=title,
|
|
content=content,
|
|
tags=tags,
|
|
),
|
|
ensure_ascii=False,
|
|
default=str,
|
|
)
|
|
|
|
|
|
@function_tool(timeout=30)
|
|
async def delete_note(ctx: RunContextWrapper, note_id: str) -> str:
|
|
"""Delete a note.
|
|
|
|
Args:
|
|
note_id: Note id to delete.
|
|
"""
|
|
return json.dumps(
|
|
await asyncio.to_thread(_delete_note_impl, note_id), ensure_ascii=False, default=str
|
|
)
|