mirror of
https://github.com/magnus919/agent-skills.git
synced 2026-09-12 03:56:53 +03:00
fix: relocate binary analysis skill
This commit is contained in:
username
@@ -0,0 +1,24 @@
|
||||
"""Optional local worker — IPC server and client.
|
||||
|
||||
The worker is an optional background process that maintains a warm backend
|
||||
adapter, reducing cold-start costs for repeated analysis operations. When the
|
||||
worker is not running, all commands function identically in one-shot mode.
|
||||
|
||||
Components:
|
||||
- WorkerServer: Unix domain socket IPC server with warm adapter
|
||||
- WorkerClient: Client for communicating with the worker
|
||||
- get_worker_status(): Convenience function to check worker state
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from binary_analysis.worker.client import WorkerClient, get_worker_status, read_pid
|
||||
from binary_analysis.worker.server import WorkerServer, run_worker
|
||||
|
||||
__all__ = [
|
||||
"WorkerClient",
|
||||
"WorkerServer",
|
||||
"get_worker_status",
|
||||
"read_pid",
|
||||
"run_worker",
|
||||
]
|
||||
@@ -0,0 +1,217 @@
|
||||
"""Worker IPC client — connects to the worker server for warm-backend requests.
|
||||
|
||||
When the worker is available, commands can route through the client for
|
||||
faster response times (avoiding cold-start costs). When the worker is
|
||||
unavailable, commands fall back to one-shot mode transparently.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import json
|
||||
import os
|
||||
import socket
|
||||
from typing import Any
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Path helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
WORKER_DIR = os.path.join(os.path.expanduser("~"), ".binary-analysis")
|
||||
|
||||
|
||||
def _socket_path() -> str:
|
||||
"""Return the path to the worker Unix domain socket."""
|
||||
return os.path.join(WORKER_DIR, "worker.sock")
|
||||
|
||||
|
||||
def _pid_path() -> str:
|
||||
"""Return the path to the worker PID file."""
|
||||
return os.path.join(WORKER_DIR, "worker.pid")
|
||||
|
||||
|
||||
def _started_at_path() -> str:
|
||||
"""Return the path to the worker started-at timestamp file."""
|
||||
return os.path.join(WORKER_DIR, "worker.started_at")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Worker client
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class WorkerClient:
|
||||
"""Client for communicating with the worker IPC server.
|
||||
|
||||
Usage::
|
||||
|
||||
client = WorkerClient()
|
||||
if client.is_available():
|
||||
result = client.send_request({"action": "execute", "command": "metadata", ...})
|
||||
# use worker-backed result
|
||||
else:
|
||||
# fall back to one-shot mode
|
||||
"""
|
||||
|
||||
def __init__(self, timeout: float = 10.0) -> None:
|
||||
self._timeout = timeout
|
||||
|
||||
def is_available(self) -> bool:
|
||||
"""Check whether the worker is running and reachable.
|
||||
|
||||
Returns True if we can connect to the worker socket and get a
|
||||
successful ping response.
|
||||
"""
|
||||
sock_path = _socket_path()
|
||||
if not os.path.exists(sock_path):
|
||||
return False
|
||||
|
||||
# Also check that the PID file is valid
|
||||
if not _is_pid_alive():
|
||||
return False
|
||||
|
||||
try:
|
||||
result = self.send_request({"action": "ping"})
|
||||
return result.get("success", False) is True
|
||||
except (OSError, ConnectionRefusedError, TimeoutError):
|
||||
return False
|
||||
|
||||
def send_request(self, request: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Send a request to the worker and return the response.
|
||||
|
||||
Args:
|
||||
request: A dict with at minimum an "action" field.
|
||||
|
||||
Returns:
|
||||
The JSON-decoded response dict.
|
||||
|
||||
Raises:
|
||||
OSError: If connection fails.
|
||||
TimeoutError: If the connection times out.
|
||||
json.JSONDecodeError: If the response is not valid JSON.
|
||||
"""
|
||||
sock_path = _socket_path()
|
||||
if not os.path.exists(sock_path):
|
||||
raise OSError(f"Worker socket not found: {sock_path}")
|
||||
|
||||
sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
||||
sock.settimeout(self._timeout)
|
||||
|
||||
try:
|
||||
sock.connect(sock_path)
|
||||
|
||||
# Send request (single JSON line)
|
||||
payload = json.dumps(request).encode("utf-8") + b"\n"
|
||||
sock.sendall(payload)
|
||||
|
||||
# Read response (single JSON line)
|
||||
response_data = b""
|
||||
while b"\n" not in response_data:
|
||||
chunk = sock.recv(65536)
|
||||
if not chunk:
|
||||
break
|
||||
response_data += chunk
|
||||
|
||||
if not response_data:
|
||||
raise OSError("Worker closed connection without response")
|
||||
|
||||
result: dict[str, Any] = json.loads(response_data.decode("utf-8").strip())
|
||||
return result
|
||||
finally:
|
||||
with contextlib.suppress(OSError):
|
||||
sock.close()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Process management helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _is_pid_alive() -> bool:
|
||||
"""Check if the PID in the PID file corresponds to a running process."""
|
||||
pid_path = _pid_path()
|
||||
if not os.path.exists(pid_path):
|
||||
return False
|
||||
|
||||
try:
|
||||
with open(pid_path) as f:
|
||||
pid_str = f.read().strip()
|
||||
if not pid_str:
|
||||
return False
|
||||
pid = int(pid_str)
|
||||
except (ValueError, OSError):
|
||||
return False
|
||||
|
||||
try:
|
||||
os.kill(pid, 0)
|
||||
return True
|
||||
except OSError:
|
||||
return False
|
||||
|
||||
|
||||
def read_pid() -> int | None:
|
||||
"""Read the worker PID from the PID file.
|
||||
|
||||
Returns None if the PID file doesn't exist, is empty, or is invalid.
|
||||
"""
|
||||
pid_path = _pid_path()
|
||||
if not os.path.exists(pid_path):
|
||||
return None
|
||||
|
||||
try:
|
||||
with open(pid_path) as f:
|
||||
pid_str = f.read().strip()
|
||||
if not pid_str:
|
||||
return None
|
||||
return int(pid_str)
|
||||
except (ValueError, OSError):
|
||||
return None
|
||||
|
||||
|
||||
def read_started_at() -> float | None:
|
||||
"""Read the worker started_at timestamp from the file.
|
||||
|
||||
Returns None if the file doesn't exist or is invalid.
|
||||
"""
|
||||
path = _started_at_path()
|
||||
if not os.path.exists(path):
|
||||
return None
|
||||
|
||||
try:
|
||||
with open(path) as f:
|
||||
value = f.read().strip()
|
||||
if not value:
|
||||
return None
|
||||
return float(value)
|
||||
except (ValueError, OSError):
|
||||
return None
|
||||
|
||||
|
||||
def get_worker_status() -> dict[str, Any]:
|
||||
"""Get the current worker status.
|
||||
|
||||
Returns a dict with:
|
||||
- state: "running" or "stopped"
|
||||
- pid: integer PID when running, null when stopped
|
||||
- uptime_seconds: float when running, null when stopped
|
||||
"""
|
||||
pid = read_pid()
|
||||
if pid is not None and _is_pid_alive():
|
||||
started_at = read_started_at()
|
||||
import time
|
||||
|
||||
uptime = None
|
||||
if started_at is not None:
|
||||
uptime = time.monotonic() - started_at
|
||||
|
||||
return {
|
||||
"state": "running",
|
||||
"pid": pid,
|
||||
"uptime_seconds": round(uptime, 3) if uptime is not None else None,
|
||||
}
|
||||
else:
|
||||
return {
|
||||
"state": "stopped",
|
||||
"pid": None,
|
||||
"uptime_seconds": None,
|
||||
}
|
||||
@@ -0,0 +1,59 @@
|
||||
"""Adapter resolution — try worker first, fall back to one-shot mode.
|
||||
|
||||
Provides a helper for CLI commands to resolve a backend adapter,
|
||||
transparently routing through the worker when available and falling
|
||||
back to direct (one-shot) initialization when the worker is not running.
|
||||
|
||||
Usage::
|
||||
|
||||
from binary_analysis.worker.resolver import resolve_adapter
|
||||
|
||||
adapter, source = resolve_adapter()
|
||||
# adapter is a FakeAdapter (or other BackendAdapter)
|
||||
# source is "worker" or "one-shot"
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from binary_analysis.adapters.fake import FakeAdapter
|
||||
|
||||
|
||||
def resolve_adapter() -> tuple[FakeAdapter, str]:
|
||||
"""Resolve a backend adapter, preferring worker when available.
|
||||
|
||||
Returns:
|
||||
A tuple of (adapter, source) where:
|
||||
- adapter: A configured FakeAdapter instance
|
||||
- source: "worker" if served by the worker, "one-shot" otherwise
|
||||
|
||||
When the worker is running, the adapter returned is a one-shot
|
||||
adapter (the worker integration is transparent to callers — the
|
||||
CLI commands already work in one-shot mode and the worker is an
|
||||
optional optimization that can be layered on later).
|
||||
"""
|
||||
from binary_analysis.worker.client import WorkerClient
|
||||
|
||||
client = WorkerClient(timeout=2.0)
|
||||
if client.is_available():
|
||||
# In the full implementation, the worker would serve the adapter.
|
||||
# For now, we fall back to one-shot but report the source.
|
||||
# The worker is an optional optimization; all commands must work
|
||||
# without it.
|
||||
pass
|
||||
|
||||
# Always use one-shot mode for now. Commands work identically
|
||||
# whether the worker is running or not.
|
||||
adapter = FakeAdapter()
|
||||
adapter.set_fixture("pe-default", FakeAdapter.pe_fixture())
|
||||
adapter.set_fixture("elf-default", FakeAdapter.elf_fixture())
|
||||
adapter.set_fixture("macho-default", FakeAdapter.macho_fixture())
|
||||
|
||||
return adapter, "one-shot"
|
||||
|
||||
|
||||
def is_worker_available() -> bool:
|
||||
"""Check if the worker is running and reachable."""
|
||||
from binary_analysis.worker.client import WorkerClient
|
||||
|
||||
client = WorkerClient(timeout=2.0)
|
||||
return client.is_available()
|
||||
@@ -0,0 +1,290 @@
|
||||
"""Worker IPC server — maintains a warm backend adapter for fast reuse.
|
||||
|
||||
The worker listens on a Unix domain socket (loopback only — no network exposure).
|
||||
It uses a simple JSON-line protocol: each request is a single JSON line,
|
||||
each response is a single JSON line.
|
||||
|
||||
The worker maintains a single FakeAdapter instance (or GhidraAdapter when
|
||||
configured) that stays warm across requests, avoiding cold-start costs.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import json
|
||||
import os
|
||||
import signal
|
||||
import socket
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
from binary_analysis.adapters.fake import FakeAdapter
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Constants
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
WORKER_DIR = os.path.join(os.path.expanduser("~"), ".binary-analysis")
|
||||
DEFAULT_BUFFER_SIZE = 65536
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# PID file helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _ensure_worker_dir() -> str:
|
||||
"""Create the worker runtime directory if it doesn't exist."""
|
||||
os.makedirs(WORKER_DIR, exist_ok=True)
|
||||
return WORKER_DIR
|
||||
|
||||
|
||||
def _pid_path() -> str:
|
||||
"""Return the path to the worker PID file."""
|
||||
return os.path.join(WORKER_DIR, "worker.pid")
|
||||
|
||||
|
||||
def _socket_path() -> str:
|
||||
"""Return the path to the worker Unix domain socket."""
|
||||
return os.path.join(WORKER_DIR, "worker.sock")
|
||||
|
||||
|
||||
def _started_at_path() -> str:
|
||||
"""Return the path to the worker started-at timestamp file."""
|
||||
return os.path.join(WORKER_DIR, "worker.started_at")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Worker server
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class WorkerServer:
|
||||
"""IPC server that maintains a warm backend adapter.
|
||||
|
||||
The server accepts connections on a Unix domain socket and processes
|
||||
JSON-line requests. Each request must include an "action" field
|
||||
("execute", "ping", or "shutdown").
|
||||
|
||||
The server runs in the foreground; daemonization is handled by the
|
||||
``binary worker start`` CLI command via fork.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._adapter: FakeAdapter | None = None
|
||||
self._running = False
|
||||
self._started_at: float = 0.0
|
||||
self._socket: socket.socket | None = None
|
||||
|
||||
@property
|
||||
def adapter(self) -> FakeAdapter:
|
||||
"""Return the warm backend adapter, initializing on first access."""
|
||||
if self._adapter is None:
|
||||
self._adapter = FakeAdapter()
|
||||
self._adapter.set_fixture("pe-default", FakeAdapter.pe_fixture())
|
||||
self._adapter.set_fixture("elf-default", FakeAdapter.elf_fixture())
|
||||
self._adapter.set_fixture("macho-default", FakeAdapter.macho_fixture())
|
||||
return self._adapter
|
||||
|
||||
@property
|
||||
def started_at(self) -> float:
|
||||
"""Return the monotonic start time of the worker."""
|
||||
return self._started_at
|
||||
|
||||
def start(self) -> None:
|
||||
"""Start the worker server.
|
||||
|
||||
Creates the PID file, socket, and starts accepting connections.
|
||||
Blocks until shutdown is requested.
|
||||
"""
|
||||
_ensure_worker_dir()
|
||||
|
||||
# Remove any stale socket
|
||||
sock_path = _socket_path()
|
||||
if os.path.exists(sock_path):
|
||||
os.unlink(sock_path)
|
||||
|
||||
# Write PID file
|
||||
pid = os.getpid()
|
||||
with open(_pid_path(), "w") as f:
|
||||
f.write(str(pid))
|
||||
|
||||
# Write started_at timestamp
|
||||
self._started_at = time.monotonic()
|
||||
with open(_started_at_path(), "w") as f:
|
||||
f.write(str(self._started_at))
|
||||
|
||||
# Pre-warm the adapter
|
||||
_ = self.adapter
|
||||
|
||||
# Create and bind socket
|
||||
server_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
||||
server_sock.bind(sock_path)
|
||||
server_sock.listen(5)
|
||||
self._socket = server_sock
|
||||
self._running = True
|
||||
|
||||
# Set up signal handlers for graceful shutdown
|
||||
signal.signal(signal.SIGTERM, self._handle_signal)
|
||||
signal.signal(signal.SIGINT, self._handle_signal)
|
||||
|
||||
while self._running:
|
||||
try:
|
||||
server_sock.settimeout(1.0)
|
||||
conn, _addr = server_sock.accept()
|
||||
self._handle_connection(conn)
|
||||
except TimeoutError:
|
||||
continue
|
||||
except OSError:
|
||||
break
|
||||
|
||||
self._cleanup()
|
||||
|
||||
def stop(self) -> None:
|
||||
"""Signal the server to stop."""
|
||||
self._running = False
|
||||
if self._socket:
|
||||
with contextlib.suppress(OSError):
|
||||
self._socket.close()
|
||||
|
||||
def _handle_signal(self, signum: int, frame: Any) -> None:
|
||||
"""Handle SIGTERM/SIGINT for graceful shutdown."""
|
||||
self.stop()
|
||||
|
||||
def _handle_connection(self, conn: socket.socket) -> None:
|
||||
"""Handle a single client connection."""
|
||||
conn.settimeout(30.0)
|
||||
data = b""
|
||||
while True:
|
||||
try:
|
||||
chunk = conn.recv(DEFAULT_BUFFER_SIZE)
|
||||
if not chunk:
|
||||
break
|
||||
data += chunk
|
||||
if b"\n" in data:
|
||||
break
|
||||
except TimeoutError:
|
||||
break
|
||||
|
||||
if data:
|
||||
# Parse request (single JSON line)
|
||||
try:
|
||||
request: dict[str, Any] = json.loads(data.decode("utf-8").strip())
|
||||
except (json.JSONDecodeError, UnicodeDecodeError):
|
||||
response: dict[str, Any] = {"success": False, "error": "Invalid JSON request"}
|
||||
conn.sendall((json.dumps(response) + "\n").encode("utf-8"))
|
||||
else:
|
||||
action = request.get("action", "")
|
||||
|
||||
if action == "ping":
|
||||
response = {"success": True, "pong": True, "pid": os.getpid()}
|
||||
elif action == "shutdown":
|
||||
response = {"success": True, "message": "Shutting down"}
|
||||
conn.sendall((json.dumps(response) + "\n").encode("utf-8"))
|
||||
self.stop()
|
||||
with contextlib.suppress(OSError):
|
||||
conn.close()
|
||||
return
|
||||
elif action == "execute":
|
||||
response = self._execute_command(request)
|
||||
else:
|
||||
response = {"success": False, "error": f"Unknown action: {action}"}
|
||||
|
||||
conn.sendall((json.dumps(response) + "\n").encode("utf-8"))
|
||||
|
||||
with contextlib.suppress(OSError):
|
||||
conn.close()
|
||||
|
||||
def _execute_command(self, request: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Execute a command through the warm backend adapter.
|
||||
|
||||
In the current version, the worker serves a subset of commands.
|
||||
For commands not yet routed through the worker, the CLI falls back
|
||||
to one-shot mode transparently.
|
||||
"""
|
||||
cmd = request.get("command", "")
|
||||
|
||||
if cmd == "metadata":
|
||||
return self._exec_metadata(request)
|
||||
else:
|
||||
return {"success": False, "error": f"Unsupported worker command: {cmd}"}
|
||||
|
||||
def _exec_metadata(self, request: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Execute a metadata request through the warm adapter."""
|
||||
project_path = request.get("project_path", "")
|
||||
|
||||
from uuid import UUID
|
||||
|
||||
from binary_analysis.domain.entities import Binary
|
||||
from binary_analysis.projects.manifest import load_manifest
|
||||
|
||||
manifest = load_manifest(project_path)
|
||||
binary_data = manifest.get("binary", {})
|
||||
raw_id = str(binary_data.get("id", ""))
|
||||
try:
|
||||
binary_uuid = UUID(raw_id) if raw_id else UUID(int=0)
|
||||
except ValueError:
|
||||
binary_uuid = UUID(int=0)
|
||||
|
||||
binary_entity = Binary(
|
||||
id=binary_uuid,
|
||||
sha256=str(binary_data.get("sha256", "")),
|
||||
path=str(binary_data.get("path", "")),
|
||||
format=str(binary_data.get("format", "unknown")),
|
||||
size_bytes=int(binary_data.get("size_bytes", 0)),
|
||||
)
|
||||
|
||||
metadata = self.adapter.get_metadata(binary_entity)
|
||||
entry_point = metadata.entry_point
|
||||
return {
|
||||
"success": True,
|
||||
"data": {
|
||||
"format": metadata.format,
|
||||
"architecture": metadata.architecture,
|
||||
"endianness": metadata.endianness,
|
||||
"size_bytes": metadata.size_bytes,
|
||||
"entry_point": (
|
||||
{
|
||||
"space": entry_point.space,
|
||||
"offset": entry_point.offset,
|
||||
"display": entry_point.display,
|
||||
}
|
||||
if entry_point
|
||||
else None
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
def _cleanup(self) -> None:
|
||||
"""Clean up PID file, socket, and other resources."""
|
||||
# Remove PID file
|
||||
pid_path = _pid_path()
|
||||
if os.path.exists(pid_path):
|
||||
with contextlib.suppress(OSError):
|
||||
os.unlink(pid_path)
|
||||
|
||||
# Remove socket
|
||||
sock_path = _socket_path()
|
||||
if os.path.exists(sock_path):
|
||||
with contextlib.suppress(OSError):
|
||||
os.unlink(sock_path)
|
||||
|
||||
# Close socket
|
||||
if self._socket:
|
||||
with contextlib.suppress(OSError):
|
||||
self._socket.close()
|
||||
|
||||
self._running = False
|
||||
|
||||
|
||||
def run_worker() -> None:
|
||||
"""Entry point for running the worker server in the foreground.
|
||||
|
||||
Used by ``binary worker start`` after forking.
|
||||
"""
|
||||
server = WorkerServer()
|
||||
server.start()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_worker()
|
||||
Reference in New Issue
Block a user