async def test_cmd_compile_all_rules_no_workspace(initialize_msg,
initialized_msg,
open_streams, test_rules,
yara_server):
''' Ensure CompileAllRules only compiles opened files when no workspace is specified '''
notifications = []
expected = {"result": None}
expected_msg = "wrong usage of identifier \"cuckoo\""
file_path = test_rules.joinpath("code_completion.yara").resolve()
contents = dedent("""
import \"cuckoo\"
rule ModuleCompletionExample { condition: cuckoo }
""")
did_change = {
"jsonrpc": "2.0",
"method": "textDocument/didChange",
"params": {
"textDocument": {
"uri": helpers.create_file_uri(file_path),
"version": 42
},
"contentChanges": [{
"text": contents
}]
}
}
execute_cmd = {
"jsonrpc": "2.0",
"id": 1,
"method": "workspace/executeCommand",
"params": {
"command": "yara.CompileAllRules",
"arguments": []
}
}
# keep non-existant test workspace
# to ensure YARA rules actually get properly identified
init_ws = json.loads(initialize_msg)
init_ws["params"]["rootPath"] = None
init_ws["params"]["rootUri"] = None
init_ws["params"]["workspaceFolders"] = []
# initialize server with the workspace we want to work from
reader, writer = open_streams
await yara_server.write_data(json.dumps(init_ws), writer)
await yara_server.read_request(reader)
await yara_server.write_data(initialized_msg, writer)
await yara_server.read_request(reader)
# notify the server of a file change
await yara_server.write_data(json.dumps(did_change), writer)
# finally, execute command
await yara_server.write_data(json.dumps(execute_cmd), writer)
response = await yara_server.read_request(reader)
while response.get("method") == "textDocument/publishDiagnostics":
notifications.append(response["params"])
response = await yara_server.read_request(reader)
assert response["result"] == expected
assert len(notifications) == 1
assert notifications[0]["uri"] == helpers.create_file_uri(file_path)
assert len(notifications[0]["diagnostics"]) == 1
assert notifications[0]["diagnostics"][0]["message"] == expected_msg
async def test_references_wildcard(test_rules, yara_server):
''' Ensure wildcard variables return references to all possible variables within rule they are found '''
peek_rules = str(test_rules.joinpath("peek_rules.yara").resolve())
file_uri = helpers.create_file_uri(peek_rules)
params = {
"textDocument": {
"uri": file_uri
},
"position": {
"line": 30,
"character": 11
},
"context": {
"includeDeclaration": True
}
}
document = yara_server._get_document(file_uri, dirty_files={})
result = await yara_server.provide_reference(params, document)
assert len(result) == 2
for index, location in enumerate(result):
assert isinstance(location, protocol.Location) is True
assert location.uri == file_uri
if index == 0:
assert location.range.start.line == 19
assert location.range.start.char == 9
assert location.range.end.line == 19
assert location.range.end.char == 19
elif index == 1:
assert location.range.start.line == 20
assert location.range.start.char == 9
assert location.range.end.line == 20
assert location.range.end.char == 20
async def test_references_rules(test_rules, yara_server):
''' Ensure references to rules are returned at the start of the rule name '''
peek_rules = str(test_rules.joinpath("peek_rules.yara").resolve())
file_uri = helpers.create_file_uri(peek_rules)
params = {
"textDocument": {
"uri": file_uri
},
"position": {
"line": 42,
"character": 12
},
"context": {
"includeDeclaration": True
}
}
document = yara_server._get_document(file_uri, dirty_files={})
result = await yara_server.provide_reference(params, document)
assert len(result) == 2
for index, location in enumerate(result):
assert isinstance(location, protocol.Location) is True
assert location.uri == file_uri
if index == 0:
assert location.range.start.line == 5
assert location.range.start.char == 5
assert location.range.end.line == 5
assert location.range.end.char == 18
elif index == 1:
assert location.range.start.line == 42
assert location.range.start.char == 8
assert location.range.end.line == 42
assert location.range.end.char == 21
async def test_hover_dirty_file(initialize_msg, initialized_msg, open_streams, test_rules, yara_server):
''' Ensure a variable's value is provided on hover for a dirty file '''
peek_rules = str(test_rules.joinpath("peek_rules.yara").resolve())
file_uri = helpers.create_file_uri(peek_rules)
unsaved_changes = "rule ResolveSymbol {\n strings:\n $a = \"test\"\n condition:\n #a > 3\n}\n"
did_change_msg = json.dumps({
"jsonrpc": "2.0", "method": "textDocument/didChange",
"params": {
"textDocument": {"uri": file_uri},
"contentChanges": [{"text": unsaved_changes}]
}
})
hover_msg = json.dumps({
"jsonrpc": "2.0", "method": "textDocument/hover", "id": 2,
"params": {
"textDocument": {"uri": file_uri},
"position": {"line": 4, "character": 3}
}
})
reader, writer = open_streams
await yara_server.write_data(initialize_msg, writer)
await yara_server.read_request(reader)
await yara_server.write_data(initialized_msg, writer)
await yara_server.read_request(reader)
await yara_server.write_data(did_change_msg, writer)
await yara_server.write_data(hover_msg, writer)
response = await yara_server.read_request(reader)
# TODO: build JSON decoder to convert JSON objects to protocol objects
assert response["result"]["contents"]["kind"] == "plaintext"
assert response["result"]["contents"]["value"] == "\"test\""
writer.close()
await writer.wait_closed()
async def test_exceptions_handled(initialize_msg, initialized_msg,
open_streams, test_rules, yara_server):
''' Ensure server notifies user when errors are encountered '''
expected = {
"jsonrpc": "2.0",
"method": "window/showMessage",
"params": {
"type": 1,
"message": "Could not find symbol for definition request"
}
}
peek_rules = str(test_rules.joinpath("peek_rules.yara").resolve())
file_uri = helpers.create_file_uri(peek_rules)
error_request = json.dumps({
"jsonrpc": "2.0",
"id": 1, # the initialize message takes id 0
"method": "textDocument/definition",
"params": {
"textDocument": {
"uri": file_uri
},
"position": {}
}
})
reader, writer = open_streams
await yara_server.write_data(initialize_msg, writer)
await yara_server.read_request(reader)
await yara_server.write_data(initialized_msg, writer)
await yara_server.read_request(reader)
await yara_server.write_data(error_request, writer)
response = await yara_server.read_request(reader)
assert response == expected
writer.close()
await writer.wait_closed()
async def test_compile_on_save_true(caplog, init_server, open_streams, test_rules, yara_server):
''' Ensure diagnostics are returned on save when 'compile_on_save' set to true '''
expected_msg = "syntax error, unexpected <true>, expecting text string"
expected_sev = protocol.DiagnosticSeverity.ERROR
new_config = {"compile_on_save": True}
peek_rules = str(test_rules.joinpath("peek_rules.yara").resolve())
file_uri = helpers.create_file_uri(peek_rules)
change_config_msg = json.dumps({
"jsonrpc":"2.0", "method": "workspace/didChangeConfiguration",
"params": {"settings": {"yara": new_config}}
})
save_file_msg = json.dumps({
"jsonrpc":"2.0", "method": "textDocument/didSave",
"params": {"textDocument": {"uri": file_uri}}
})
reader, writer = open_streams
with caplog.at_level(logging.DEBUG, "yara"):
await init_server(reader, writer, yara_server)
await yara_server.write_data(change_config_msg, writer)
await yara_server.write_data(save_file_msg, writer)
response = await yara_server.read_request(reader)
diagnostics = response["params"]["diagnostics"]
assert len(diagnostics) == 1
assert diagnostics[0]["message"] == expected_msg
assert diagnostics[0]["severity"] == expected_sev
assert ("yara", logging.DEBUG, "Changed workspace config to {}".format(json.dumps(new_config))) in caplog.record_tuples
writer.close()
await writer.wait_closed()
async def test_format_alt_tabsize(test_rules, yara_server):
''' Ensure a text edit is provided on format with tabSize set '''
expected = dedent("""\
rule Oneline : test
{
strings:
$a = "test"
condition:
$a
}""")
oneline = str(test_rules.joinpath("oneline.yar").resolve())
file_uri = helpers.create_file_uri(oneline)
message = {
"params": {
"textDocument": {
"uri": file_uri
},
"position": {
"line": 29,
"character": 12
},
"options": {
"tabSize": 2
}
}
}
result = await yara_server.provide_formatting(message, True)
assert len(result) == 1
edit = result[0]
assert isinstance(edit, protocol.TextEdit) is True
assert edit.newText == expected
async def test_renames(test_rules, yara_server):
''' Ensure variables can be renamed '''
peek_rules = str(test_rules.joinpath("peek_rules.yara").resolve())
file_uri = helpers.create_file_uri(peek_rules)
# @dstring[1]: Line 30, Col 12
new_text = "test_rename"
params = {
"textDocument": {
"uri": file_uri
},
"position": {
"line": 29,
"character": 12
},
"newName": new_text
}
document = yara_server._get_document(file_uri, dirty_files={})
result = await yara_server.provide_rename(params, document, file_uri)
assert isinstance(result, protocol.WorkspaceEdit) is True
assert len(result.changes) == 3
acceptable_lines = [21, 28, 29]
for edit in result.changes:
assert isinstance(edit, protocol.TextEdit) is True
assert edit.newText == new_text
assert edit.range.start.line in acceptable_lines
def test_create_file_uri(test_rules):
''' Ensure file URIs are generated from paths '''
test_rule_path = str(test_rules)
expected = "file://{}".format(
quote(test_rule_path.replace("\\", "/"), safe="/\\"))
output = helpers.create_file_uri(test_rule_path)
assert output == expected
async def test_dirty_files(test_rules, yara_server):
''' Ensure server prefers versions of dirty files over those backed by file path '''
peek_rules = str(test_rules.joinpath("peek_rules.yara").resolve())
file_uri = helpers.create_file_uri(peek_rules)
unsaved_changes = "rule ResolveSymbol {\n strings:\n $a = \"test\"\n condition:\n #a > 3\n}\n"
dirty_files = {file_uri: unsaved_changes}
document = yara_server._get_document(file_uri, dirty_files)
assert document == unsaved_changes
async def test_no_hover(test_rules, yara_server):
''' Ensure non-variables do not return hovers '''
peek_rules = str(test_rules.joinpath("peek_rules.yara").resolve())
file_uri = helpers.create_file_uri(peek_rules)
message = {
"params": {
"textDocument": {"uri": file_uri},
"position": {"line": 25, "character": 12}
}
}
result = await yara_server.provide_hover(message, True)
assert result is None
async def test_code_completion_unexpected(test_rules, yara_server):
''' Ensure code completion doesn't return items or error out when a symbol does not have any items to be completed '''
code_completion = str(test_rules.joinpath("code_completion.yara").resolve())
file_uri = helpers.create_file_uri(code_completion)
message = {
"params": {
"textDocument": {"uri": file_uri},
"position": {"line": 8, "character": 25}
}
}
result = await yara_server.provide_code_completion(message, True)
assert result == []
async def test_hover(test_rules, yara_server):
''' Ensure a variable's value is provided on hover '''
peek_rules = str(test_rules.joinpath("peek_rules.yara").resolve())
file_uri = helpers.create_file_uri(peek_rules)
message = {
"params": {
"textDocument": {"uri": file_uri},
"position": {"line": 29, "character": 12}
}
}
result = await yara_server.provide_hover(message, True)
assert isinstance(result, protocol.Hover) is True
assert result.contents.kind == protocol.MarkupKind.Plaintext
assert result.contents.value == "\"double string\" wide nocase fullword"
async def test_format_keep_whitespace(test_rules, yara_server):
''' Ensure a text edit is provided with untrimmed whitespace '''
expected = dedent("""\
rule Oneline : test
{
strings:
$a = "test"
condition:
$a
}""")
oneline = str(test_rules.joinpath("oneline.yar").resolve())
file_uri = helpers.create_file_uri(oneline)
# spacing should be preserved
dirty_files = {file_uri: expected}
file_uri = helpers.create_file_uri(oneline)
message = {
"params": {
"textDocument": {
"uri": file_uri
},
"position": {
"line": 29,
"character": 12
},
"options": {
"trimTrailingWhitespace": False
}
}
}
result = await yara_server.provide_formatting(message,
True,
dirty_files=dirty_files)
assert len(result) == 1
edit = result[0]
assert isinstance(edit, protocol.TextEdit) is True
assert edit.newText == expected
async def test_code_completion_module_function(test_rules, yara_server):
''' Ensure code completion returns a properly-formatted snippet string when a module entry is a function '''
expected = [
protocol.CompletionItem("is_dll", protocol.CompletionItemKind.FUNCTION, detail="pe.is_dll()", insertText="is_dll()")
]
code_completion = str(test_rules.joinpath("code_completion.yara").resolve())
file_uri = helpers.create_file_uri(code_completion)
message = {
"params": {
"textDocument": {"uri": file_uri},
"position": {"line": 26, "character": 17}
}
}
result = await yara_server.provide_code_completion(message, True)
assert len(result) == len(expected)
assert result == expected
async def test_no_hover(test_rules, yara_server):
''' Ensure non-variables do not return hovers '''
peek_rules = str(test_rules.joinpath("peek_rules.yara").resolve())
file_uri = helpers.create_file_uri(peek_rules)
params = {
"textDocument": {
"uri": file_uri
},
"position": {
"line": 25,
"character": 12
}
}
document = yara_server._get_document(file_uri, dirty_files={})
result = await yara_server.provide_hover(params, document)
assert result is None
async def test_no_references(test_rules, yara_server):
''' Ensure server does not return references if none are found '''
alienspy = str(test_rules.joinpath("apt_alienspy_rat.yar").resolve())
file_uri = helpers.create_file_uri(alienspy)
params = {
"textDocument": {
"uri": file_uri
},
"position": {
"line": 47,
"character": 99
},
}
document = yara_server._get_document(file_uri, dirty_files={})
result = await yara_server.provide_reference(params, document)
assert result == []
async def test_code_completion_overflow(test_rules, yara_server):
''' Ensure code completion doesn't return items or error out when a position doesn't exist in the file '''
code_completion = str(
test_rules.joinpath("code_completion.yara").resolve())
file_uri = helpers.create_file_uri(code_completion)
params = {
"textDocument": {
"uri": file_uri
},
"position": {
"line": 9,
"character": 25
}
}
document = yara_server._get_document(file_uri, dirty_files={})
result = await yara_server.provide_code_completion(params, document)
assert result == []
async def test_code_completion_regular(test_rules, yara_server):
''' Ensure code completion works with functions defined in modules schema '''
code_completion = str(test_rules.joinpath("code_completion.yara").resolve())
expected = [
protocol.CompletionItem("network", protocol.CompletionItemKind.CLASS, detail="cuckoo.network"),
protocol.CompletionItem("registry", protocol.CompletionItemKind.CLASS, detail="cuckoo.registry"),
protocol.CompletionItem("filesystem", protocol.CompletionItemKind.CLASS, detail="cuckoo.filesystem"),
protocol.CompletionItem("sync", protocol.CompletionItemKind.CLASS, detail="cuckoo.sync")
]
file_uri = helpers.create_file_uri(code_completion)
message = {
"params": {
"textDocument": {"uri": file_uri},
"position": {"line": 10, "character": 15}
}
}
actual = await yara_server.provide_code_completion(message, True)
assert len(actual) == len(expected)
assert actual == expected
async def test_no_definitions(test_rules, yara_server):
''' Ensure no definition is provided for symbols that are not variables or rules '''
peek_rules = str(test_rules.joinpath("peek_rules.yara").resolve())
file_uri = helpers.create_file_uri(peek_rules)
params = {
"textDocument": {
"uri": file_uri
},
"position": {
"line": 27,
"character": 12
},
"context": {
"includeDeclaration": True
}
}
document = yara_server._get_document(file_uri, dirty_files={})
result = await yara_server.provide_definition(params, document)
assert result == []
async def test_format_notify_user(test_rules, uninstall_pkg, yara_server):
''' Ensure the formatter notifies the user if plyara is not installed '''
expected_msg = "plyara is not installed. Formatting is disabled"
oneline = str(test_rules.joinpath("oneline.yar").resolve())
file_uri = helpers.create_file_uri(oneline)
message = {
"params": {
"textDocument": {
"uri": file_uri
},
"position": {
"line": 29,
"character": 12
}
}
}
await uninstall_pkg("plyara")
with pytest.raises(ce.NoDependencyFound) as excinfo:
await yara_server.provide_formatting(message, True)
assert expected_msg == str(excinfo.value)
async def test_format_no_imports(test_rules, yara_server):
''' Ensure imports are removed from provided rules. They should not be affected by formatter '''
rulefile = str(test_rules.joinpath("code_completion.yara").resolve())
file_uri = helpers.create_file_uri(rulefile)
message = {
"params": {
"textDocument": {
"uri": file_uri
},
"position": {
"line": 9,
"character": 12
}
}
}
result = await yara_server.provide_formatting(message, True)
assert len(result) == 3
assert all([isinstance(edit, protocol.TextEdit) for edit in result])
full_text = "\n".join([edit.newText for edit in result])
# should only be two imports - one for cuckoo and one for pe
assert full_text.count("import ") == 0
async def test_cancel(initialize_msg, initialized_msg, open_streams,
test_rules, yara_server):
''' Ensure a task can be cancelled before it returns '''
# initialize the server
reader, writer = open_streams
await yara_server.write_data(initialize_msg, writer)
await yara_server.read_request(reader)
await yara_server.write_data(initialized_msg, writer)
await yara_server.read_request(reader)
# execute a task that shouldn't complete on its own
alienspy = str(test_rules.joinpath("apt_alienspy_rat.yar").resolve())
file_uri = helpers.create_file_uri(alienspy)
msg_id = 1
message = {
"jsonrpc": "2.0",
"id": msg_id,
"method": "textDocument/hover",
"params": {
"textDocument": {
"uri": file_uri
},
"position": {
"line": 38,
"character": 15
}
}
}
await yara_server.write_data(json.dumps(message), writer)
await yara_server.read_request(reader)
# cancel that task
cancel = {
"jsonrpc": "2.0",
"method": "$/cancelRequest",
"params": {
"id": msg_id
}
}
await yara_server.write_data(json.dumps(cancel), writer)
response = await yara_server.read_request(reader)
print(response)
async def test_definitions_variables_regular(test_rules, yara_server):
''' Ensure definition is provided for a normal variable '''
peek_rules = str(test_rules.joinpath("peek_rules.yara").resolve())
file_uri = helpers.create_file_uri(peek_rules)
params = {
"textDocument": {
"uri": file_uri
},
"position": {
"line": 24,
"character": 12
}
}
document = yara_server._get_document(file_uri, dirty_files={})
result = await yara_server.provide_definition(params, document)
assert len(result) == 1
assert isinstance(result[0], protocol.Location) is True
assert result[0].uri == file_uri
assert result[0].range.start.line == 19
assert result[0].range.start.char == 9
assert result[0].range.end.line == 19
assert result[0].range.end.char == 22
async def test_format_keep_newlines(test_rules, yara_server):
''' Ensure a text edit is provided with extra newlines '''
expected = dedent("""\
rule Oneline : test
{
strings:
$a = "test"
condition:
$a
}
""")
oneline = str(test_rules.joinpath("oneline.yar").resolve())
with open(oneline) as ifile:
file_uri = helpers.create_file_uri(oneline)
dirty_files = {file_uri: "%s\n\n\n" % ifile.read()}
message = {
"params": {
"textDocument": {
"uri": file_uri
},
"position": {
"line": 29,
"character": 12
},
"options": {
"trimFinalNewlines": False
}
}
}
result = await yara_server.provide_formatting(message,
True,
dirty_files=dirty_files)
assert len(result) == 1
edit = result[0]
assert isinstance(edit, protocol.TextEdit) is True
assert edit.newText == expected
async def test_cmd_compile_all_rules(initialize_msg, initialized_msg,
open_streams, test_rules, yara_server):
''' Ensure CompileAllRules compiles all YARA rule files in the given workspace using the "executeCommand" action '''
expected = {"result": None}
execute_cmd = {
"jsonrpc": "2.0",
"id": 1,
"method": "workspace/executeCommand",
"params": {
"command": "yara.CompileAllRules",
"arguments": []
}
}
# change workspace from non-existant test workspace to our test_rules
# to ensure YARA rules actually get properly identified
init_ws = json.loads(initialize_msg)
init_ws["params"]["rootPath"] = str(test_rules)
init_ws["params"]["rootUri"] = helpers.create_file_uri(str(test_rules))
init_ws["params"]["workspaceFolders"] = [{
"uri":
init_ws["params"]["rootUri"],
"name":
test_rules.name
}]
# initialize server with the workspace we want to work from
reader, writer = open_streams
await yara_server.write_data(json.dumps(init_ws), writer)
await yara_server.read_request(reader)
await yara_server.write_data(initialized_msg, writer)
await yara_server.read_request(reader)
# finally, execute command
await yara_server.write_data(json.dumps(execute_cmd), writer)
response = await yara_server.read_request(reader)
# diagnostics are sent as notifications, and there may be an arbitrary number depending on test files
while response.get("method") == "textDocument/publishDiagnostics":
response = await yara_server.read_request(reader)
# once all the diagnostics are finished, we should get our final response, which should be null
assert response.get("result", {}) == expected
async def test_code_completion_module_dictionary(test_rules, yara_server):
''' Ensure code completion returns a properly-formatted list of strings when a module entry is a list of dictionary keys '''
actual = []
expected = []
options = [
"Comments", "CompanyName", "FileDescription", "FileVersion", "InternalName",
"LegalCopyright", "LegalTrademarks", "OriginalFilename", "ProductName", "ProductVersion"
]
for option in options:
snippet = "version_info[\"{}\"]".format(option)
detail = "pe.{}".format(snippet)
expected.append(protocol.CompletionItem(option, protocol.CompletionItemKind.INTERFACE, detail=detail, insertText=snippet))
code_completion = str(test_rules.joinpath("code_completion.yara").resolve())
file_uri = helpers.create_file_uri(code_completion)
message = {
"params": {
"textDocument": {"uri": file_uri},
"position": {"line": 18, "character": 15}
}
}
actual = await yara_server.provide_code_completion(message, True)
assert len(actual) == len(expected)
assert actual == expected
async def test_definitions_private_rules(test_rules, yara_server):
''' Ensure definition is provided for a private rule name '''
private_goto_rules = str(
test_rules.joinpath("private_rule_goto.yara").resolve())
file_uri = helpers.create_file_uri(private_goto_rules)
params = {
"textDocument": {
"uri": file_uri
},
"position": {
"line": 9,
"character": 14
}
}
document = yara_server._get_document(file_uri, dirty_files={})
result = await yara_server.provide_definition(params, document)
assert len(result) == 1
assert isinstance(result[0], protocol.Location) is True
assert result[0].uri == file_uri
assert result[0].range.start.line == 0
assert result[0].range.start.char == 13
assert result[0].range.end.line == 0
assert result[0].range.end.char == 28
async def test_definitions_variables_length(test_rules, yara_server):
''' Ensure definition is provided for a variable with length modifier (!) '''
peek_rules = str(test_rules.joinpath("peek_rules.yara").resolve())
file_uri = helpers.create_file_uri(peek_rules)
message = {
"params": {
"textDocument": {
"uri": file_uri
},
"position": {
"line": 42,
"character": 32
}
}
}
result = await yara_server.provide_definition(message, True)
assert len(result) == 1
assert isinstance(result[0], protocol.Location) is True
assert result[0].uri == file_uri
assert result[0].range.start.line == 40
assert result[0].range.start.char == 9
assert result[0].range.end.line == 40
assert result[0].range.end.char == 22
async def test_code_completion_regular(test_rules, yara_server):
''' Ensure code completion works with functions defined in modules schema '''
actual = []
code_completion = str(
test_rules.joinpath("code_completion.yara").resolve())
expected = ["network", "registry", "filesystem", "sync"]
file_uri = helpers.create_file_uri(code_completion)
params = {
"textDocument": {
"uri": file_uri
},
"position": {
"line": 9,
"character": 15
}
}
document = yara_server._get_document(file_uri, dirty_files={})
result = await yara_server.provide_code_completion(params, document)
assert len(result) == 4
for completion in result:
assert isinstance(completion, protocol.CompletionItem) is True
actual.append(completion.label)
assert actual == expected
