315 lines
14 KiB
Python
315 lines
14 KiB
Python
#!/usr/bin/env python3
|
|
import base64
|
|
import os
|
|
import tempfile
|
|
from typing import Optional, Dict, Any
|
|
|
|
from yams_sdk.base import BasePlugin
|
|
from yams_sdk.decorators import rpc
|
|
|
|
|
|
class GhidraPlugin(BasePlugin):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self._project_dir: Optional[str] = None
|
|
self._project_name: str = "YamsGhidra"
|
|
self._ghidra_install: Optional[str] = os.environ.get("GHIDRA_INSTALL_DIR")
|
|
self._started = False
|
|
# Register RPCs discovered via decorator metadata
|
|
for name in dir(self):
|
|
fn = getattr(self, name)
|
|
rpc_name = getattr(fn, "__rpc_name__", None)
|
|
if rpc_name:
|
|
self.register(rpc_name, fn)
|
|
|
|
def manifest(self) -> dict:
|
|
return {"name": "yams_ghidra", "version": "0.0.1", "interfaces": ["ghidra_analysis_v1"]}
|
|
|
|
def init(self, config: Dict[str, Any]) -> None: # noqa: A003
|
|
# Lazy-import pyghidra on init to avoid hard dependency for non-users
|
|
try:
|
|
import pyghidra # type: ignore
|
|
except Exception as e: # noqa: BLE001
|
|
raise RuntimeError(f"pyghidra not available: {e}")
|
|
|
|
self._ghidra_install = config.get("ghidra_install", self._ghidra_install)
|
|
self._project_dir = config.get("project_dir") or tempfile.mkdtemp(prefix="yams-ghidra-")
|
|
self._project_name = config.get("project_name", self._project_name)
|
|
|
|
if not self._started:
|
|
# Start JVM and initialize Ghidra in headless mode
|
|
if self._ghidra_install:
|
|
pyghidra.start(install_dir=self._ghidra_install)
|
|
else:
|
|
pyghidra.start()
|
|
self._started = True
|
|
|
|
# Probe project creation/open (no persistent handle kept; opened per call)
|
|
try:
|
|
with pyghidra.open_project(self._project_dir, self._project_name, create=True):
|
|
pass
|
|
except Exception as e: # noqa: BLE001
|
|
raise RuntimeError(f"Failed to open/create Ghidra project: {e}")
|
|
|
|
def health(self) -> dict:
|
|
return {"status": "ok", "started": self._started, "project_dir": self._project_dir}
|
|
|
|
@staticmethod
|
|
def _materialize_source(source: Dict[str, Any]) -> str:
|
|
st = source.get("type")
|
|
if st == "path":
|
|
return source["path"]
|
|
if st == "bytes":
|
|
data_b64 = source.get("data")
|
|
if not data_b64:
|
|
raise ValueError("bytes source requires 'data' base64 field")
|
|
data = base64.b64decode(data_b64)
|
|
fd, p = tempfile.mkstemp(prefix="yams-ghidra-bin-")
|
|
os.write(fd, data)
|
|
os.close(fd)
|
|
return p
|
|
raise ValueError("unsupported source.type; expected 'path' or 'bytes'")
|
|
|
|
def _open_program_ctx(self, project, path):
|
|
import pyghidra # type: ignore
|
|
try:
|
|
# Preferred convenience (may be deprecated in future)
|
|
return pyghidra.open_program(path, analyze=True) # type: ignore[attr-defined]
|
|
except AttributeError:
|
|
# Fallback: load via ProgramLoader then access via program_context
|
|
loader = pyghidra.program_loader().project(project).source(path)
|
|
with loader.load() as load_results:
|
|
load_results.save(pyghidra.dummy_monitor())
|
|
return pyghidra.program_context(project, "/" + os.path.basename(path))
|
|
|
|
def _iter_functions(self, program, limit: int, offset: int = 0):
|
|
listing = program.getListing()
|
|
it = listing.getFunctions(True)
|
|
skipped = 0
|
|
collected = 0
|
|
while it.hasNext():
|
|
f = it.next()
|
|
if skipped < offset:
|
|
skipped += 1
|
|
continue
|
|
yield f
|
|
collected += 1
|
|
if collected >= limit:
|
|
break
|
|
|
|
@rpc("ghidra.analyze")
|
|
def analyze(self, source: Dict[str, Any], opts: Optional[Dict[str, Any]] = None) -> dict:
|
|
import pyghidra # type: ignore
|
|
path = self._materialize_source(source)
|
|
max_functions = int((opts or {}).get("max_functions", 50))
|
|
with pyghidra.open_project(self._project_dir, self._project_name, create=True) as project:
|
|
with self._open_program_ctx(project, path) as ctx:
|
|
# ctx may be FlatProgramAPI or Program; normalize
|
|
try:
|
|
program = ctx.getCurrentProgram()
|
|
except AttributeError:
|
|
program = ctx
|
|
lang = program.getLanguage().getLanguageID().getIdAsString()
|
|
funcs = []
|
|
for f in self._iter_functions(program, limit=max_functions):
|
|
funcs.append({"name": f.getName(), "addr": f.getEntryPoint().toString()})
|
|
return {"arch": lang, "count": len(funcs), "functions": funcs}
|
|
|
|
@rpc("ghidra.list_functions")
|
|
def list_functions(self, source: Dict[str, Any], opts: Optional[Dict[str, Any]] = None) -> dict:
|
|
import pyghidra # type: ignore
|
|
path = self._materialize_source(source)
|
|
limit = int((opts or {}).get("limit", 100))
|
|
offset = int((opts or {}).get("offset", 0))
|
|
with pyghidra.open_project(self._project_dir, self._project_name, create=True) as project:
|
|
with self._open_program_ctx(project, path) as ctx:
|
|
try:
|
|
program = ctx.getCurrentProgram()
|
|
except AttributeError:
|
|
program = ctx
|
|
items = []
|
|
for f in self._iter_functions(program, limit=limit, offset=offset):
|
|
items.append({"name": f.getName(), "addr": f.getEntryPoint().toString()})
|
|
# total is best-effort; iterate fully if requested
|
|
total = offset + len(items)
|
|
return {"items": items, "total": total}
|
|
|
|
@rpc("ghidra.get_function")
|
|
def get_function(self, source: Dict[str, Any], func: Dict[str, Any], include_body: bool = False) -> dict:
|
|
import pyghidra # type: ignore
|
|
path = self._materialize_source(source)
|
|
with pyghidra.open_project(self._project_dir, self._project_name, create=True) as project:
|
|
with self._open_program_ctx(project, path) as ctx:
|
|
try:
|
|
program = ctx.getCurrentProgram()
|
|
except AttributeError:
|
|
program = ctx
|
|
listing = program.getListing()
|
|
addr_factory = program.getAddressFactory()
|
|
target = None
|
|
if "addr" in func:
|
|
a = addr_factory.getAddress(func["addr"])
|
|
target = listing.getFunctionAt(a)
|
|
elif "name" in func:
|
|
it = listing.getFunctions(True)
|
|
while it.hasNext():
|
|
f = it.next()
|
|
if f.getName() == func["name"]:
|
|
target = f
|
|
break
|
|
if target is None:
|
|
return {"ok": False, "error": "NotFound"}
|
|
out = {"name": target.getName(), "addr": target.getEntryPoint().toString()}
|
|
if include_body:
|
|
# Best-effort: dump code units around entry (placeholder)
|
|
out["body"] = f"Function {target.getName()} at {out['addr']}"
|
|
return out
|
|
|
|
@rpc("ghidra.decompile_function")
|
|
def decompile_function(self, source: Dict[str, Any], func: Dict[str, Any], opts: Optional[Dict[str, Any]] = None) -> dict:
|
|
import pyghidra # type: ignore
|
|
try:
|
|
from ghidra.app.decompiler import DecompInterface # type: ignore
|
|
except Exception:
|
|
return {"ok": False, "error": "DecompilerUnavailable"}
|
|
path = self._materialize_source(source)
|
|
with pyghidra.open_project(self._project_dir, self._project_name, create=True) as project:
|
|
with self._open_program_ctx(project, path) as ctx:
|
|
try:
|
|
program = ctx.getCurrentProgram()
|
|
except AttributeError:
|
|
program = ctx
|
|
listing = program.getListing()
|
|
addr_factory = program.getAddressFactory()
|
|
target = None
|
|
if "addr" in func:
|
|
a = addr_factory.getAddress(func["addr"])
|
|
target = listing.getFunctionAt(a)
|
|
elif "name" in func:
|
|
it = listing.getFunctions(True)
|
|
while it.hasNext():
|
|
f = it.next()
|
|
if f.getName() == func["name"]:
|
|
target = f
|
|
break
|
|
if target is None:
|
|
return {"ok": False, "error": "NotFound"}
|
|
di = DecompInterface()
|
|
if not di.openProgram(program):
|
|
return {"ok": False, "error": "OpenProgramFailed"}
|
|
res = di.decompileFunction(target, int((opts or {}).get("timeout_sec", 30)), pyghidra.dummy_monitor())
|
|
if not res or not res.getDecompiledFunction():
|
|
return {"ok": False, "error": "DecompileFailed"}
|
|
decomp = res.getDecompiledFunction().getC()
|
|
return {"ok": True, "decomp": decomp, "meta": {"name": target.getName(), "addr": target.getEntryPoint().toString()}}
|
|
|
|
@rpc("ghidra.search")
|
|
def search(self, source: Dict[str, Any], query: Dict[str, Any], opts: Optional[Dict[str, Any]] = None) -> dict:
|
|
"""Simple search across functions by name or decompiled text substring.
|
|
query: { text: "..." }
|
|
opts: { max_hits?: int }
|
|
"""
|
|
import pyghidra # type: ignore
|
|
text = (query or {}).get("text", "").strip()
|
|
if not text:
|
|
return {"hits": [], "total": 0}
|
|
path = self._materialize_source(source)
|
|
max_hits = int((opts or {}).get("max_hits", 25))
|
|
hits = []
|
|
total = 0
|
|
try:
|
|
from ghidra.app.decompiler import DecompInterface # type: ignore
|
|
except Exception:
|
|
DecompInterface = None
|
|
with pyghidra.open_project(self._project_dir, self._project_name, create=True) as project:
|
|
with self._open_program_ctx(project, path) as ctx:
|
|
try:
|
|
program = ctx.getCurrentProgram()
|
|
except AttributeError:
|
|
program = ctx
|
|
listing = program.getListing()
|
|
it = listing.getFunctions(True)
|
|
di = None
|
|
if DecompInterface is not None:
|
|
di = DecompInterface()
|
|
di.openProgram(program)
|
|
while it.hasNext():
|
|
f = it.next()
|
|
name = f.getName()
|
|
addr = f.getEntryPoint().toString()
|
|
found = text.lower() in name.lower()
|
|
snippet = ""
|
|
if not found and di is not None:
|
|
try:
|
|
res = di.decompileFunction(f, int((opts or {}).get("timeout_sec", 10)), pyghidra.dummy_monitor())
|
|
if res and res.getDecompiledFunction():
|
|
c = res.getDecompiledFunction().getC()
|
|
if text.lower() in c.lower():
|
|
found = True
|
|
# create a small snippet around first occurrence
|
|
idx = c.lower().find(text.lower())
|
|
start = max(0, idx - 40)
|
|
end = min(len(c), idx + 40)
|
|
snippet = c[start:end]
|
|
except Exception:
|
|
pass
|
|
if found:
|
|
total += 1
|
|
if len(hits) < max_hits:
|
|
hits.append({"title": name, "addr": addr, "snippet": snippet})
|
|
return {"hits": hits, "total": total}
|
|
|
|
@rpc("ghidra.grep")
|
|
def grep(self, source: Dict[str, Any], pattern: Dict[str, Any], opts: Optional[Dict[str, Any]] = None) -> dict:
|
|
"""Regex grep over decompiled function text.
|
|
pattern: { regex: "...", flags?: "i" }
|
|
opts: { max_hits?: int }
|
|
"""
|
|
import pyghidra # type: ignore
|
|
import re
|
|
regex = (pattern or {}).get("regex", "")
|
|
if not regex:
|
|
return {"matches": [], "total": 0}
|
|
flags = 0
|
|
if (pattern or {}).get("flags") == "i":
|
|
flags |= re.IGNORECASE
|
|
rx = re.compile(regex, flags)
|
|
path = self._materialize_source(source)
|
|
max_hits = int((opts or {}).get("max_hits", 25))
|
|
matches = []
|
|
total = 0
|
|
try:
|
|
from ghidra.app.decompiler import DecompInterface # type: ignore
|
|
except Exception:
|
|
return {"matches": [], "total": 0, "error": "DecompilerUnavailable"}
|
|
with pyghidra.open_project(self._project_dir, self._project_name, create=True) as project:
|
|
with self._open_program_ctx(project, path) as ctx:
|
|
try:
|
|
program = ctx.getCurrentProgram()
|
|
except AttributeError:
|
|
program = ctx
|
|
listing = program.getListing()
|
|
it = listing.getFunctions(True)
|
|
di = DecompInterface()
|
|
if not di.openProgram(program):
|
|
return {"matches": [], "total": 0, "error": "OpenProgramFailed"}
|
|
while it.hasNext():
|
|
f = it.next()
|
|
try:
|
|
res = di.decompileFunction(f, int((opts or {}).get("timeout_sec", 10)), pyghidra.dummy_monitor())
|
|
if not res or not res.getDecompiledFunction():
|
|
continue
|
|
c = res.getDecompiledFunction().getC()
|
|
for m in rx.finditer(c):
|
|
total += 1
|
|
if len(matches) < max_hits:
|
|
start = max(0, m.start() - 40)
|
|
end = min(len(c), m.end() + 40)
|
|
matches.append({"path": f.getName(), "addr": f.getEntryPoint().toString(), "snippet": c[start:end]})
|
|
except Exception:
|
|
continue
|
|
return {"matches": matches, "total": total}
|
|
|
|
|
|
if __name__ == "__main__":
|
|
GhidraPlugin().run()
|