reverse-engineering-assistant 2.9.2__tar.gz → 2.9.4__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (73) hide show
  1. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/PKG-INFO +3 -1
  2. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/pyproject.toml +3 -1
  3. reverse_engineering_assistant-2.9.4/reverse_engineering_assistant/api_server_tools/__init__.py +91 -0
  4. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/api_server_tools/connection.py +3 -0
  5. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/api_server_tools/llm_tools.py +22 -47
  6. reverse_engineering_assistant-2.9.4/reverse_engineering_assistant/api_server_tools/re_tool_box/__init__.py +6 -0
  7. reverse_engineering_assistant-2.9.4/reverse_engineering_assistant/api_server_tools/re_tool_box/bookmarks.py +65 -0
  8. reverse_engineering_assistant-2.9.4/reverse_engineering_assistant/api_server_tools/re_tool_box/comment.py +65 -0
  9. reverse_engineering_assistant-2.9.4/reverse_engineering_assistant/api_server_tools/re_tool_box/cross_reference.py +69 -0
  10. reverse_engineering_assistant-2.9.4/reverse_engineering_assistant/api_server_tools/re_tool_box/cursor.py +45 -0
  11. reverse_engineering_assistant-2.9.4/reverse_engineering_assistant/api_server_tools/re_tool_box/data.py +154 -0
  12. reverse_engineering_assistant-2.9.4/reverse_engineering_assistant/api_server_tools/re_tool_box/decompilation.py +204 -0
  13. reverse_engineering_assistant-2.9.4/reverse_engineering_assistant/api_server_tools/re_tool_box/symbols.py +209 -0
  14. reverse_engineering_assistant-2.9.4/reverse_engineering_assistant/api_server_tools/re_tools.py +17 -0
  15. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/assistant.py +192 -110
  16. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/assistant_api_server.py +9 -6
  17. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/chat_client.py +106 -19
  18. reverse_engineering_assistant-2.9.4/reverse_engineering_assistant/crash_dump.py +112 -0
  19. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/model.py +13 -9
  20. reverse_engineering_assistant-2.9.4/reverse_engineering_assistant/protocol/RevaBookmark_pb2.py +35 -0
  21. reverse_engineering_assistant-2.9.4/reverse_engineering_assistant/protocol/RevaBookmark_pb2.pyi +33 -0
  22. reverse_engineering_assistant-2.9.4/reverse_engineering_assistant/protocol/RevaBookmark_pb2_grpc.py +145 -0
  23. reverse_engineering_assistant-2.9.4/reverse_engineering_assistant/protocol/RevaChat_pb2.py +39 -0
  24. reverse_engineering_assistant-2.9.4/reverse_engineering_assistant/protocol/RevaChat_pb2.pyi +55 -0
  25. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/protocol/RevaChat_pb2_grpc.py +46 -1
  26. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/protocol/RevaComment_pb2.py +6 -6
  27. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/protocol/RevaComment_pb2.pyi +6 -4
  28. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/protocol/RevaComment_pb2_grpc.py +2 -1
  29. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/protocol/RevaData_pb2.py +11 -3
  30. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/protocol/RevaData_pb2.pyi +22 -0
  31. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/protocol/RevaData_pb2_grpc.py +45 -1
  32. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/protocol/RevaGetCursor_pb2.py +4 -4
  33. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/protocol/RevaGetCursor_pb2.pyi +6 -4
  34. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/protocol/RevaGetCursor_pb2_grpc.py +2 -1
  35. reverse_engineering_assistant-2.9.4/reverse_engineering_assistant/protocol/RevaGetDecompilation_pb2.py +44 -0
  36. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/protocol/RevaGetDecompilation_pb2.pyi +16 -2
  37. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/protocol/RevaGetDecompilation_pb2_grpc.py +45 -1
  38. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/protocol/RevaGetSymbols_pb2.py +14 -14
  39. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/protocol/RevaGetSymbols_pb2.pyi +6 -4
  40. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/protocol/RevaGetSymbols_pb2_grpc.py +3 -1
  41. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/protocol/RevaHandshake_pb2_grpc.py +2 -1
  42. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/protocol/RevaHeartbeat_pb2_grpc.py +2 -1
  43. reverse_engineering_assistant-2.9.4/reverse_engineering_assistant/protocol/RevaReferences_pb2.py +31 -0
  44. reverse_engineering_assistant-2.9.4/reverse_engineering_assistant/protocol/RevaReferences_pb2.pyi +20 -0
  45. reverse_engineering_assistant-2.9.2/reverse_engineering_assistant/protocol/RevaGetReferences_pb2_grpc.py → reverse_engineering_assistant-2.9.4/reverse_engineering_assistant/protocol/RevaReferences_pb2_grpc.py +17 -16
  46. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/protocol/RevaVariable_pb2_grpc.py +1 -1
  47. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant.egg-info/PKG-INFO +3 -1
  48. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant.egg-info/SOURCES.txt +15 -3
  49. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant.egg-info/requires.txt +2 -0
  50. reverse_engineering_assistant-2.9.2/reverse_engineering_assistant/api_server_tools/__init__.py +0 -16
  51. reverse_engineering_assistant-2.9.2/reverse_engineering_assistant/api_server_tools/re_tools.py +0 -504
  52. reverse_engineering_assistant-2.9.2/reverse_engineering_assistant/protocol/RevaChat_pb2.py +0 -31
  53. reverse_engineering_assistant-2.9.2/reverse_engineering_assistant/protocol/RevaChat_pb2.pyi +0 -27
  54. reverse_engineering_assistant-2.9.2/reverse_engineering_assistant/protocol/RevaGetDecompilation_pb2.py +0 -40
  55. reverse_engineering_assistant-2.9.2/reverse_engineering_assistant/protocol/RevaGetReferences_pb2.py +0 -31
  56. reverse_engineering_assistant-2.9.2/reverse_engineering_assistant/protocol/RevaGetReferences_pb2.pyi +0 -20
  57. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/README.md +0 -0
  58. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/.gitignore +0 -0
  59. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/__init__.py +0 -0
  60. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/documents.py +0 -0
  61. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/protocol/RevaHandshake_pb2.py +0 -0
  62. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/protocol/RevaHandshake_pb2.pyi +0 -0
  63. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/protocol/RevaHeartbeat_pb2.py +0 -0
  64. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/protocol/RevaHeartbeat_pb2.pyi +0 -0
  65. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/protocol/RevaVariable_pb2.py +0 -0
  66. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/protocol/RevaVariable_pb2.pyi +0 -0
  67. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/protocol/__init__.py +0 -0
  68. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/reva_exceptions.py +0 -0
  69. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant/tool.py +0 -0
  70. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant.egg-info/dependency_links.txt +0 -0
  71. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant.egg-info/entry_points.txt +0 -0
  72. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/reverse_engineering_assistant.egg-info/top_level.txt +0 -0
  73. {reverse_engineering_assistant-2.9.2 → reverse_engineering_assistant-2.9.4}/setup.cfg +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: reverse-engineering-assistant
3
- Version: 2.9.2
3
+ Version: 2.9.4
4
4
  Summary: An AI assistant for reverse engineering tasks
5
5
  Author: サイバーカイダ (cyberkaida)
6
6
  Classifier: License :: OSI Approved :: Apache Software License
@@ -12,6 +12,8 @@ Requires-Dist: langchain
12
12
  Requires-Dist: langchain-core
13
13
  Requires-Dist: langchain-openai
14
14
  Requires-Dist: langchain-community
15
+ Requires-Dist: langchain-experimental
16
+ Requires-Dist: langgraph
15
17
  Requires-Dist: prompt_toolkit
16
18
  Requires-Dist: sentence_transformers
17
19
  Requires-Dist: PyYAML
@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
5
5
  [project]
6
6
  name = "reverse-engineering-assistant"
7
7
  readme = "README.md"
8
- version = "2.9.2"
8
+ version = "2.9.4"
9
9
  authors = [
10
10
  {name="サイバーカイダ (cyberkaida)"},
11
11
  ]
@@ -21,6 +21,8 @@ dependencies = [
21
21
  "langchain-core",
22
22
  "langchain-openai",
23
23
  "langchain-community",
24
+ "langchain-experimental",
25
+ "langgraph",
24
26
  "prompt_toolkit",
25
27
  "sentence_transformers",
26
28
  "PyYAML",
@@ -0,0 +1,91 @@
1
+
2
+
3
+ from typing import List
4
+ from reverse_engineering_assistant.reva_exceptions import RevaToolException
5
+ from reverse_engineering_assistant.tool import AssistantProject
6
+ from reverse_engineering_assistant.assistant import RevaTool
7
+ from reverse_engineering_assistant.model import RevaModel
8
+ from reverse_engineering_assistant.api_server_tools.connection import get_channel
9
+ from typing import Optional, Tuple
10
+ import logging
11
+
12
+
13
+ class RevaMessageHandler(object):
14
+ handles_type = None
15
+
16
+ _global_message_handlers: List[RevaMessageHandler] = []
17
+
18
+ __all__ = ['register_message_handler']
19
+
20
+ def register_message_handler(cls: RevaMessageHandler):
21
+ _global_message_handlers.append(cls)
22
+ return cls
23
+
24
+
25
+ class RevaRemoteTool(RevaTool):
26
+ logger: logging.Logger
27
+
28
+ def __init__(self, project: AssistantProject, llm: RevaModel) -> None:
29
+ self.logger = logging.getLogger(f"reverse_engineering_assistant.RevaRemoteTool.{self.__class__.__name__}")
30
+ self.logger.addHandler(logging.FileHandler(project.project_path / "reva.log"))
31
+ super().__init__(project, llm)
32
+
33
+ @property
34
+ def channel(self):
35
+ return get_channel()
36
+
37
+ def resolve_to_address_and_symbol(self, thing: str) -> Tuple[str, Optional[str]]:
38
+ """
39
+ Resolve a string to an address and symbol.
40
+ If it is an address it can be a namespaced address or a plain hex address.
41
+
42
+ This helps reduce hallucinations and catch issues with symbols and namespaces early.
43
+
44
+ Returns a tuple of (address, symbol).
45
+ """
46
+ self.logger.debug(f"Resolving {thing} to address and symbol")
47
+ assert thing is not None
48
+ address: Optional[str] = None
49
+ symbol: Optional[str] = None
50
+ try:
51
+ # This is a plain address in the main namespace
52
+ address = hex(int(thing, 16))
53
+ self.logger.debug(f"Resolved {thing} to address: {address}")
54
+ except ValueError:
55
+ # This could also be a Ghidra namespaced address, so let's check that too!
56
+ if "::" in thing:
57
+ last_part = thing.split("::")[-1]
58
+ try:
59
+ hex(int(last_part, 16))
60
+ # If the last part of the address is a hex number, then we can assume it is an address
61
+ address = thing
62
+ except ValueError:
63
+ # Otherwise, it is a symbol
64
+ symbol = thing
65
+ else:
66
+ # This is a symbol
67
+ symbol = thing
68
+
69
+ # We can check if it is a symbol
70
+ from ..protocol import RevaGetSymbols_pb2_grpc, RevaGetSymbols_pb2
71
+ stub = RevaGetSymbols_pb2_grpc.RevaToolSymbolServiceStub(self.channel)
72
+ request = RevaGetSymbols_pb2.RevaSymbolRequest()
73
+ if address:
74
+ request.address = address
75
+ if symbol:
76
+ request.name = symbol
77
+
78
+ self.logger.debug(f"Getting symbol for {thing} request: {request}")
79
+ response = stub.GetSymbol(request)
80
+ self.logger.debug(f"Got symbol for {thing} response: {response}")
81
+
82
+ if response.name:
83
+ symbol = response.name
84
+ if response.address:
85
+ address = response.address
86
+
87
+ if address is None and symbol is None:
88
+ raise RevaToolException(message=f"Could not resolve {thing} to an address or symbol. Double check your symbol or address is correct.")
89
+ self.logger.debug(f"Resolved {thing} to address: {address}, symbol: {symbol}")
90
+ assert address is not None
91
+ return address, symbol
@@ -11,6 +11,9 @@ from functools import cache
11
11
  _channel: Optional[Channel] = None
12
12
  @cache
13
13
  def get_channel() -> Channel:
14
+ """
15
+ Get the global gRPC channel to the extension.
16
+ """
14
17
  global _channel
15
18
  if not _channel:
16
19
  raise ValueError("Channel not set")
@@ -4,6 +4,7 @@ from ast import Call
4
4
  import queue
5
5
  import threading
6
6
  from typing import Callable, Dict, Optional
7
+ from uuid import uuid4
7
8
  from venv import logger
8
9
 
9
10
  from ..protocol.RevaChat_pb2_grpc import RevaChatServiceServicer
@@ -18,58 +19,27 @@ module_logger = logging.getLogger("reva-server")
18
19
  from langchain_core.callbacks.base import BaseCallbackHandler
19
20
  from langchain_core.agents import AgentAction, AgentFinish
20
21
 
21
- from langchain_core.language_models.base import BaseLanguageModel
22
- from langchain_core.language_models.chat_models import BaseChatModel
23
-
24
- class RevaActionCollector(BaseCallbackHandler):
25
- """
26
- A callback handler for logging agent actions in the reverse engineering assistant.
27
-
28
- This class logs agent actions and calls a callback. This is what prints the green
29
- thoughts from the model to the console. This is very useful for the analyst to understand
30
- what the model is doing (and is arguably the most important part of the assistant output!)
31
-
32
- Attributes:
33
- callback (Callable[[str], None]): The callback function to call when an agent action is performed.
34
- logger (logging.Logger): The logger instance for the reverse_engineering_assistant.RevaActionLogger class.
35
- """
36
-
37
- callback: Callable[[str], None]
38
- def __init__(self, callback: Callable[[str], None]) -> None:
39
- super().__init__()
40
- self.callback = callback
41
-
42
- logger = logging.getLogger("reverse_engineering_assistant.RevaActionLogger")
43
-
44
- def on_agent_action(self, action: AgentAction, **kwargs) -> None:
45
- """
46
- Callback method called when an agent action is performed.
47
-
48
- Args:
49
- action (AgentAction): The agent action that was performed.
50
- **kwargs: Additional keyword arguments.
51
-
52
- Returns:
53
- None
54
- """
55
- logger.debug(f"Agent action: {action} {kwargs}")
56
- # TODO: Should this be AgentAction?
57
- # TODO: Is `.action` still a thing?
58
- self.callback(str(action.log))
59
-
22
+ from reverse_engineering_assistant.model import RevaModel
23
+ from reverse_engineering_assistant.model import get_llm_ollama, get_llm_openai
60
24
 
61
25
  class RevaChat(RevaChatServiceServicer):
62
26
  logger = logging.getLogger("reva-server.RevaChat")
63
- llm: BaseLanguageModel | BaseChatModel
64
27
 
65
- def __init__(self, llm: BaseLanguageModel | BaseChatModel) -> None:
66
- self.llm = llm
28
+ def _model_from_request(self, request) -> RevaModel:
29
+ """
30
+ Given a request, return the model associated with the request.
31
+ """
32
+ if request.ollama.model:
33
+ return get_llm_ollama(base_url=request.ollama.url, model=request.ollama.model)
34
+ if request.openai.model:
35
+ return get_llm_openai(model=request.openai.model, api_key=request.openai.token)
36
+ raise ValueError("No model specified in request. Please file a bug.")
67
37
 
68
38
  def chat(self, request, context):
69
39
  self.logger.info(f"Received request: {request}")
70
40
  assistant = ReverseEngineeringAssistant(
71
41
  request.project,
72
- model=self.llm
42
+ model=self._model_from_request(request)
73
43
  )
74
44
  self.logger.info(f"Assistant: {assistant}")
75
45
  llm_response = assistant.query(request.message)
@@ -95,8 +65,8 @@ class RevaChat(RevaChatServiceServicer):
95
65
 
96
66
  assistant = ReverseEngineeringAssistant(
97
67
  request.project,
98
- model=self.llm,
99
- langchain_callbacks=[RevaActionCollector(callback)]
68
+ model=self._model_from_request(request),
69
+ logging_callbacks=[callback],
100
70
  )
101
71
  self.logger.info(f"Assistant: {assistant}")
102
72
 
@@ -133,8 +103,8 @@ class RevaChat(RevaChatServiceServicer):
133
103
  if not assistant:
134
104
  assistant = ReverseEngineeringAssistant(
135
105
  request.project,
136
- model=self.llm,
137
- langchain_callbacks=[RevaActionCollector(callback)]
106
+ model=self._model_from_request(request),
107
+ logging_callbacks=[callback],
138
108
  )
139
109
  self.logger.info(f"Received request: {request}")
140
110
  def run_query(query: str):
@@ -157,3 +127,8 @@ class RevaChat(RevaChatServiceServicer):
157
127
  done = True
158
128
  yield response
159
129
  t.join()
130
+
131
+ def shutdown(self, request, context):
132
+ self.logger.warning("Shutting down")
133
+ import sys
134
+ sys.exit(0)
@@ -0,0 +1,6 @@
1
+ """
2
+ To make sure your tool is loaded, you should import it here.
3
+ `assistant_api_server.py` imports everything from `re_tools.py`,
4
+ and `re_tools.py` imports everything from here.
5
+ """
6
+
@@ -0,0 +1,65 @@
1
+ from typing import Dict, List
2
+ import logging
3
+
4
+ import grpc
5
+
6
+ from reverse_engineering_assistant.tool import AssistantProject
7
+ from reverse_engineering_assistant.assistant import AssistantProject, register_tool
8
+ from reverse_engineering_assistant.reva_exceptions import RevaToolException
9
+ from reverse_engineering_assistant.api_server_tools import RevaRemoteTool
10
+ from langchain_core.language_models.base import BaseLanguageModel
11
+ from reverse_engineering_assistant.protocol import RevaBookmark_pb2, RevaBookmark_pb2_grpc
12
+
13
+ from reverse_engineering_assistant.model import RevaModel
14
+
15
+
16
+ @register_tool
17
+ class RevaBookmarks(RevaRemoteTool):
18
+
19
+ def __init__(self, project: AssistantProject, llm: RevaModel) -> None:
20
+ super().__init__(project, llm)
21
+ self.description = "Used for managing bookmarks in the program"
22
+
23
+ self.tool_functions = [
24
+ self.get_bookmarks,
25
+ self.add_bookmark,
26
+ ]
27
+
28
+ def get_bookmarks(self) -> List[Dict[str, str]]:
29
+ """
30
+ Return a list of Ghidra bookmarks in the program.
31
+ Use this to keep track of important locations in the program the user
32
+ has marked.
33
+ """
34
+ stub = RevaBookmark_pb2_grpc.RevaBookmarkStub(self.channel)
35
+
36
+ request = RevaBookmark_pb2.RevaGetBookmarksRequest()
37
+
38
+ bookmarks: List[Dict[str, str]] = []
39
+ for bookmark in stub.get_bookmarks(request):
40
+ bookmarks.append({
41
+ "address": bookmark.address,
42
+ "category": bookmark.category,
43
+ "description": bookmark.description,
44
+ })
45
+ return bookmarks
46
+
47
+ def add_bookmark(self, address_or_symbol: str, category: str, description: str) -> str:
48
+ """
49
+ Add a Ghidra bookmark at the given address or symbol with the given category and description.
50
+ If the category does not exist, it will be created. Use a category to group bookmarks together.
51
+ Make sure your category is descriptive and useful to the user.
52
+
53
+ Use this tool to keep note of anything important enough you want to look at it
54
+ again.
55
+ """
56
+ stub = RevaBookmark_pb2_grpc.RevaBookmarkStub(self.channel)
57
+
58
+ request = RevaBookmark_pb2.RevaAddBookmarkRequest()
59
+ request.category = f"ReVa.{category}"
60
+ request.description = description
61
+ request.address, _ = self.resolve_to_address_and_symbol(address_or_symbol)
62
+
63
+ response = stub.add_bookmark(request)
64
+
65
+ return "Added bookmark successfully"
@@ -0,0 +1,65 @@
1
+ from typing import Dict, List
2
+ import logging
3
+
4
+ import grpc
5
+
6
+ from reverse_engineering_assistant.tool import AssistantProject
7
+ from reverse_engineering_assistant.assistant import AssistantProject, register_tool
8
+ from reverse_engineering_assistant.reva_exceptions import RevaToolException
9
+ from reverse_engineering_assistant.api_server_tools import RevaRemoteTool
10
+
11
+ from reverse_engineering_assistant.protocol import RevaGetDecompilation_pb2_grpc, RevaGetDecompilation_pb2
12
+ from reverse_engineering_assistant.protocol import RevaComment_pb2_grpc, RevaComment_pb2
13
+
14
+ from reverse_engineering_assistant.model import RevaModel
15
+
16
+ @register_tool
17
+ class RevaSetComment(RevaRemoteTool):
18
+ """
19
+ A tool for setting comments on addresses, functions and symbols.
20
+ """
21
+
22
+ def __init__(self, project: AssistantProject, llm: RevaModel) -> None:
23
+ super().__init__(project, llm)
24
+ self.description = "Used for setting comments on addresses, functions and symbols"
25
+
26
+ self.tool_functions = [
27
+ self.set_comment,
28
+ self.set_multiple_comments,
29
+ ]
30
+
31
+ def set_multiple_comments(self, comments: Dict[str, str]) -> List[str]:
32
+ """
33
+ Set multiple comments at the same time.
34
+ Keys are addresses or symbols, values are the comments to set at that location.
35
+ This is more efficient than calling set_comment multiple times.
36
+ """
37
+ outputs: List[str] = []
38
+
39
+ for address_or_symbol, comment in comments.items():
40
+ outputs.append(self.set_comment(address_or_symbol=address_or_symbol, comment=comment))
41
+
42
+ return outputs
43
+
44
+ def set_comment(self, address_or_symbol: str, comment: str) -> str:
45
+ """
46
+ Set the comment at the given address, function or symbol to `comment`.
47
+ Use this when you want to add an explanation or note to a specific part
48
+ of the code.
49
+ """
50
+ stub = RevaComment_pb2_grpc.RevaCommentServiceStub(self.channel)
51
+
52
+ request = RevaComment_pb2.RevaSetCommentRequest()
53
+ request.comment = comment
54
+ address, symbol = self.resolve_to_address_and_symbol(address_or_symbol)
55
+
56
+ if address:
57
+ request.address = address
58
+ if symbol:
59
+ request.symbol = symbol
60
+
61
+ try:
62
+ response: RevaComment_pb2.RevaSetCommentResponse = stub.SetComment(request)
63
+ except grpc.RpcError as e:
64
+ raise RevaToolException(f"Failed to set comment: {e}")
65
+ return "Set comment successfully"
@@ -0,0 +1,69 @@
1
+ from typing import Dict, List, Union
2
+ import logging
3
+
4
+ import grpc
5
+
6
+ from reverse_engineering_assistant.tool import AssistantProject
7
+ from reverse_engineering_assistant.assistant import AssistantProject, register_tool
8
+ from reverse_engineering_assistant.reva_exceptions import RevaToolException
9
+ from reverse_engineering_assistant.api_server_tools import RevaRemoteTool
10
+ from reverse_engineering_assistant.model import RevaModel
11
+
12
+ @register_tool
13
+ class RevaCrossReferenceTool(RevaRemoteTool):
14
+ """
15
+ An tool to retrieve cross references, to and from, addresses.
16
+ """
17
+ def __init__(self, project: AssistantProject, llm: RevaModel) -> None:
18
+ super().__init__(project, llm)
19
+ self.description = "Used for retrieving cross references to and from addresses"
20
+
21
+ self.tool_functions = [
22
+ self.get_references,
23
+ self.get_references_to,
24
+ self.get_references_from,
25
+ ]
26
+
27
+ def get_references(self, address_or_symbol: str) -> Dict[str, Union[str, List[str]]]:
28
+ """
29
+ Return a list of references to and from the given address or symbol.
30
+ These might be calls from/to other functions, or data references from/to this address.
31
+ """
32
+ from reverse_engineering_assistant.protocol import RevaReferences_pb2_grpc, RevaReferences_pb2
33
+
34
+ stub = RevaReferences_pb2_grpc.RevaReferenceServiceStub(self.channel)
35
+
36
+ request = RevaReferences_pb2.RevaGetReferencesRequest()
37
+
38
+ address, symbol = self.resolve_to_address_and_symbol(address_or_symbol)
39
+ request.address = address
40
+
41
+ try:
42
+ response: RevaReferences_pb2.RevaGetReferencesResponse = stub.get_references(request)
43
+ except grpc.RpcError as e:
44
+ raise RevaToolException(f"Failed to get references: {e}")
45
+
46
+ result: Dict[str, Union[str, List[str]]] = {
47
+ "address": address,
48
+ "incoming_references": list(response.incoming_references),
49
+ "outgoing_references": list(response.outgoing_references),
50
+ }
51
+
52
+ if symbol:
53
+ result["symbol"] = symbol
54
+
55
+ return result
56
+
57
+ def get_references_to(self, address_or_symbol: str) -> List[str]:
58
+ """
59
+ Return a list of references to the given address or symbol.
60
+ """
61
+ references = self.get_references(address_or_symbol)
62
+ return references.get("incoming_references", []) # type: ignore
63
+
64
+ def get_references_from(self, address_or_symbol: str) -> List[str]:
65
+ """
66
+ Return a list of references from the given address or symbol.
67
+ """
68
+ references = self.get_references(address_or_symbol)
69
+ return references.get("outgoing_references", []) # type: ignore
@@ -0,0 +1,45 @@
1
+ from typing import Dict, List, Union
2
+ import logging
3
+
4
+ import grpc
5
+
6
+ from reverse_engineering_assistant.tool import AssistantProject
7
+ from reverse_engineering_assistant.assistant import AssistantProject, register_tool
8
+ from reverse_engineering_assistant.reva_exceptions import RevaToolException
9
+ from reverse_engineering_assistant.api_server_tools import RevaRemoteTool
10
+ from reverse_engineering_assistant.model import RevaModel
11
+ from reverse_engineering_assistant.protocol import RevaGetCursor_pb2, RevaGetCursor_pb2_grpc
12
+
13
+
14
+ @register_tool
15
+ class RevaGetCursor(RevaRemoteTool):
16
+
17
+ def __init__(self, project: AssistantProject, llm: RevaModel) -> None:
18
+ super().__init__(project, llm)
19
+ self.description = "Used for getting and setting the cursor or when the user mentions 'this'"
20
+
21
+ self.tool_functions = [
22
+ self.get_cursor,
23
+ ]
24
+
25
+ def get_cursor(self) -> Dict[str, Union[str, int]]:
26
+ """
27
+ Return the current location the user is looking at in the program.
28
+ Use this to find the current function, symbol or address. When the user mentions "this",
29
+ you should find the current location using this function.
30
+
31
+ This method returns a dictionary with the keys "address", "symbol", and "function".
32
+ Use other tools to gather context around this location. For example, decompile
33
+ the function to find the exact code at this location in the listing or decompilation.
34
+ """
35
+ stub = RevaGetCursor_pb2_grpc.RevaGetCursorStub(self.channel)
36
+
37
+ request = RevaGetCursor_pb2.RevaGetCursorRequest()
38
+
39
+ response = stub.getCursor(request)
40
+
41
+ return {
42
+ "address": response.address,
43
+ "symbol": response.symbol,
44
+ "function": response.function,
45
+ }
@@ -0,0 +1,154 @@
1
+ from typing import Dict, List, Optional, Union
2
+ from binascii import a2b_hex, b2a_hex, a2b_base64
3
+ import logging
4
+
5
+ import grpc
6
+
7
+ from reverse_engineering_assistant.tool import AssistantProject
8
+ from reverse_engineering_assistant.assistant import AssistantProject, register_tool
9
+ from reverse_engineering_assistant.reva_exceptions import RevaToolException
10
+ from reverse_engineering_assistant.api_server_tools import RevaRemoteTool
11
+ from reverse_engineering_assistant.model import RevaModel
12
+ from reverse_engineering_assistant.protocol import RevaGetDecompilation_pb2_grpc, RevaGetDecompilation_pb2
13
+ from reverse_engineering_assistant.protocol import RevaData_pb2_grpc, RevaData_pb2, RevaGetSymbols_pb2, RevaGetSymbols_pb2_grpc
14
+
15
+
16
+ @register_tool
17
+ class RevaData(RevaRemoteTool):
18
+ """
19
+ A tool for getting and setting data
20
+ """
21
+
22
+ def __init__(self, project: AssistantProject, llm: RevaModel) -> None:
23
+ super().__init__(project, llm)
24
+ self.description = "Used for getting and setting data"
25
+
26
+ self.tool_functions = [
27
+ self.list_strings,
28
+ self.list_data,
29
+ self.get_data,
30
+ self.set_global_data_type,
31
+ self.xor_data,
32
+ self.base64_decode_string,
33
+ ]
34
+
35
+ def list_strings(self) -> List[Dict[str, Union[str, List[str]]]]:
36
+ """
37
+ Return a list of defined strings in the program.
38
+ """
39
+ stub = RevaData_pb2_grpc.RevaDataServiceStub(self.channel)
40
+
41
+ request = RevaData_pb2.RevaStringListRequest()
42
+
43
+ defined_strings: List[Dict[str, Union[str, List[str]]]] = []
44
+ for string in stub.getStringList(request):
45
+ defined_strings.append({
46
+ "address": string.address,
47
+ "symbol": string.symbol,
48
+ "value": string.value,
49
+ "incoming_references": list(string.incoming_references),
50
+ "outgoing_references": list(string.outgoing_references),
51
+ })
52
+
53
+ return defined_strings
54
+
55
+ def list_data(self) -> List[Dict[str, Union[str, List[str]]]]:
56
+ """
57
+ Return a list of defined data in the program.
58
+ This is not all data, only the data that has been defined in the Ghidra database.
59
+
60
+ This function returns a list of dictionaries with the keys "address", "symbol", "type", "size", "incoming_references", and "outgoing_references".
61
+ """
62
+
63
+ stub = RevaData_pb2_grpc.RevaDataServiceStub(self.channel)
64
+
65
+ request = RevaData_pb2.RevaDataListRequest()
66
+
67
+ defined_data: List[Dict[str, Union[str, List[str]]]] = []
68
+ for data in stub.getListData(request):
69
+ defined_data.append({
70
+ "address": data.address,
71
+ "symbol": data.symbol,
72
+ "type": data.type,
73
+ "size": data.size,
74
+ "incoming_references": list(data.incoming_references),
75
+ "outgoing_references": list(data.outgoing_references),
76
+ })
77
+
78
+ return defined_data
79
+
80
+ def get_data(self, address_or_symbol: str, size: Optional[int] = None) -> Dict[str, Optional[Union[str, List[str], int]]]:
81
+ """
82
+ Return information about the data at the given address or with the given symbol.
83
+ Don't set the size if you don't know it, Ghidra will try to figure it out.
84
+
85
+ This function returns a dictionary with the keys "address", "symbol", "type", "size", "data", "incoming_references", and "outgoing_references".
86
+ """
87
+ stub = RevaData_pb2_grpc.RevaDataServiceStub(self.channel)
88
+ request = RevaData_pb2.RevaGetDataAtAddressRequest()
89
+
90
+ request.address, symbol = self.resolve_to_address_and_symbol(address_or_symbol)
91
+
92
+ if size:
93
+ request.size = size
94
+
95
+ response = stub.getDataAtAddress(request)
96
+
97
+ return {
98
+ "address": response.address,
99
+ "symbol": symbol,
100
+ "type": response.type,
101
+ "size": len(response.data),
102
+ "data": b2a_hex(response.data).decode("utf-8"),
103
+ "incoming_references": list(response.incoming_references),
104
+ "outgoing_references": list(response.outgoing_references),
105
+ }
106
+
107
+ def set_global_data_type(self, address_or_symbol: str, data_type: Optional[str], new_name: Optional[str]) -> str:
108
+ """
109
+ Set the data type or name of the data at the given address or symbol to the given data type.
110
+
111
+ `new_type` must be a string that can be passed to the "Set Data Type" dialog in Ghidra.
112
+ Something like `int`, `char`, `string`, `long`, `unsigned int`, `char[0x10]` or `char[16]`, or `int*` should work,
113
+ but you can use custom types from the program too.
114
+
115
+ You can't define a _new_ data type here, only use existing ones.
116
+ """
117
+ if not data_type and not new_name:
118
+ raise RevaToolException("You must provide either a data type or a new name")
119
+ address, symbol = self.resolve_to_address_and_symbol(address_or_symbol)
120
+
121
+ if data_type:
122
+ stub = RevaData_pb2_grpc.RevaDataServiceStub(self.channel)
123
+ request = RevaData_pb2.RevaSetGlobalDataTypeRequest()
124
+ request.address = address
125
+ request.data_type = data_type
126
+ response = stub.setGlobalDataType(request)
127
+
128
+ if new_name:
129
+ symbol_stub = RevaGetSymbols_pb2_grpc.RevaToolSymbolServiceStub(self.channel)
130
+ symbol_request = RevaGetSymbols_pb2.RevaSetSymbolNameRequest()
131
+ if symbol:
132
+ symbol_request.old_name = symbol
133
+ symbol_request.new_name = new_name
134
+ symbol_request.old_address = address
135
+ symbol_response = symbol_stub.SetSymbolName(symbol_request)
136
+
137
+ # TODO: Should we return truth to reduce hallucinations?
138
+ return f"{address_or_symbol} has been updated successfully"
139
+
140
+ def xor_data(self, hex_string: str, key_byte: str) -> str:
141
+ """
142
+ Given a data buffer and a key buffer, return the result of XORing the data with the key.
143
+ Returns a hex string.
144
+ """
145
+ # TODO: Should we expose cyberchef recipes?
146
+ data = bytes.fromhex(hex_string)
147
+ key = int(key_byte, 16)
148
+ return b2a_hex(bytes([byte ^ key for byte in data])).decode("utf-8")
149
+
150
+ def base64_decode_string(self, base64_string: str) -> str:
151
+ """
152
+ Given a base64 encoded string, return the decoded string.
153
+ """
154
+ return a2b_base64(base64_string).decode("utf-8")