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