-
Notifications
You must be signed in to change notification settings - Fork 13
/
completer.py
240 lines (184 loc) · 8.31 KB
/
completer.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
from __future__ import annotations
import logging
import re
from typing import Any
import lsprotocol.types as lsp_types
from pygls.workspace import Document
from yls import completion
from yls import utils
from yls.completion import CONDITION_KEYWORDS
from yls.plugin_manager_provider import PluginManagerProvider
from yls.strings import estimate_string_type
from yls.strings import string_modifiers_completion_items
log = logging.getLogger(__name__)
class Completer:
REGEXP_META = re.compile(r'^\s*(\w+)\s+=\s+"(.*)$')
REGEXP_IMPORT = re.compile(r'^import\s+"(.*)$')
def __init__(self, ls: Any):
self.ls = ls
self.completion_cache = completion.CompletionCache.from_yaramod(self.ls.ymod)
def complete(self, params: lsp_types.CompletionParams) -> lsp_types.CompletionList:
return lsp_types.CompletionList(is_incomplete=False, items=self._complete(params))
def signature_help(self, params: lsp_types.CompletionParams) -> lsp_types.SignatureHelp | None:
signatures = self._signature_help(params)
if signatures is None:
return None
return lsp_types.SignatureHelp(
signatures=signatures, active_parameter=0, active_signature=0
)
def _signature_help(
self, params: lsp_types.CompletionParams
) -> list[lsp_types.SignatureInformation] | None:
text_doc = self.ls.workspace.get_document(params.text_document.uri)
word = utils.cursor_word(text_doc, params.position)
log.debug(f'[SIGNATURE_HELP] Cursor word is "{word}"')
if word is None:
log.debug("[SIGNATURE_HELP] Aborting signature help request")
return None
# Try to extract the function call from the word
word = re.sub(r"\(.*", "", word)
word = word.replace("(", "").replace(")", "")
log.debug(f'[SIGNATURE_HELP] Normalized cursor word is "{word}"')
if not word:
log.debug(
"[SIGNATURE_HELP] Aborting signature help request - invalid word after normalization"
)
return None
symbol = self.completion_cache.get_function(word)
if not symbol:
log.debug("[SIGNATURE_HELP] Aborting signature help request - no symbol found")
return None
info: list[lsp_types.SignatureInformation] = []
# NOTE: Handler other overloads
info.extend(symbol.to_signature_information(word))
return info
def _complete(self, params: lsp_types.CompletionParams) -> list[lsp_types.CompletionItem]:
document = self.ls.workspace.get_document(params.text_document.uri)
res = []
# If the cursor is not on a word, we still want to provide completions
# for root module
word = utils.cursor_word(document, params.position) or ""
# Import completion
log.debug("[COMPLETION] Adding import completion")
res += self.complete_import(document, params.position)
# String modifiers completion
log.debug("[COMPLETION] Adding string modifiers completion")
res += self.complete_string_modifiers(document, params.position, word)
# Condition keywords completion
log.debug("[COMPLETION] Adding condition keywords completion")
res += self.complete_condition_keywords(document, params.position, word)
# Function completion
log.debug("[COMPLETION] Adding module completion")
res += self.complete_word(document, params.position, word)
# Symbols from last yara file completion
log.debug("[COMPLETION] Adding last valid yara file")
res += self.complete_last_valid_yara_file(document, params, word)
# Plugin completion
log.debug("COMPLETION] Adding completion items from plugings")
res += utils.flatten_list(
PluginManagerProvider.instance().hook.yls_completion(params=params, document=document)
)
return res
def complete_last_valid_yara_file(
self, document: Document, params: lsp_types.CompletionParams, word: str
) -> list[lsp_types.CompletionItem]:
"""Return list of completion items from last valid YaraFile."""
if self.ls.last_valid_yara_file is None:
return []
if not utils.is_in_yara_section(document, params.position.line, "condition"):
return []
res = []
# Add rules
log.debug("[COMPLETION] Adding symbols from last valid YARA file")
res += [
lsp_types.CompletionItem(label=rule.name, insert_text=rule.name, sort_text="zb")
for rule in self.ls.last_valid_yara_file.rules
if rule.name.startswith(word)
]
rule = self.ls.get_current_rule(
params.text_document.uri, params.position, self.ls.last_valid_yara_file
)
if rule is None:
return res
# Add current rule strings
res += [
lsp_types.CompletionItem(
label=string.identifier, insert_text=string.identifier, sort_text="za"
)
for string in rule.strings
if string.identifier.startswith(word)
]
return res
def complete_word(
self, document: Document, position: lsp_types.Position, word: str
) -> list[lsp_types.CompletionItem]:
"""Complete a function."""
log.debug(f'[COMPLETION] Cursor word is "{word}"')
if not utils.is_in_yara_section(document, position.line, "condition"):
return []
res = []
symbols = self.completion_cache.get_symbols_matching(word)
for symbol in symbols:
res.extend(symbol.to_completion_items())
# symbol_doc = symbol_doc.strip().replace("\t", "")
# symbol_doc = re.sub(r" {2,}", "", symbol_doc)
return res
def complete_import(
self, document: Document, position: lsp_types.Position
) -> list[lsp_types.CompletionItem]:
"""Complete an import statement."""
res = []
line = utils.cursor_line(document, position)
match_res = self.REGEXP_IMPORT.search(line)
if match_res:
log.debug("[COMPLETION] Completing import values")
for module in self.completion_cache.modules.attributes.keys():
res.append(
lsp_types.CompletionItem(
label=module,
insert_text=module,
sort_text="aaa",
kind=lsp_types.CompletionItemKind.Module,
)
)
return res
def complete_string_modifiers(
self, document: Document, position: lsp_types.Position, word: str
) -> list[lsp_types.CompletionItem]:
"""Complete string modifiers."""
# Check if the line is not empty
line = utils.cursor_line(document, position)
if not line.strip():
return []
# Complete only if end of the string is before the cursor
characters_after_cursor = set(line[position.character :])
if {'"', "/", "}"}.intersection(characters_after_cursor):
return []
# Complete only in `strings:` section
if not utils.is_in_yara_section(document, position.line, "strings"):
return []
# Filter out the modifiers based on the string type
estimated_string_type = estimate_string_type(line)
final_list = iter(string_modifiers_completion_items(estimated_string_type))
# Filter the modifiers based on the word under the cursor
final_list = filter(lambda item: item.label.startswith(word), final_list)
return list(final_list)
def complete_condition_keywords(
self, document: Document, position: lsp_types.Position, word: str
) -> list[lsp_types.CompletionItem]:
# Complete only in `condition:` section
if not utils.is_in_yara_section(document, position.line, "condition"):
return []
res = []
for keyword in CONDITION_KEYWORDS:
if not keyword.startswith(word):
continue
item = lsp_types.CompletionItem(
label=keyword,
kind=lsp_types.CompletionItemKind.Keyword,
insert_text=keyword,
insert_text_format=lsp_types.InsertTextFormat.PlainText,
sort_text="a",
)
res.append(item)
return res