diff --git a/editors/vscode/package.json b/editors/vscode/package.json index c8ec3c8..31f440f 100644 --- a/editors/vscode/package.json +++ b/editors/vscode/package.json @@ -46,6 +46,38 @@ "type": "string", "default": "", "description": "Specifies an absolute path to the directory used to store samples and module data. Used only for local YARI. (example: /home/user/samples)" + }, + "yls.snippets.metaEntries": { + "type": "object", + "default": {}, + "description": "A set of metadata entries to insert into rules. Empty values will create meta keys with a tabstop, and built-in variables are accepted (though transforms are ignored).", + "maxProperties": 50, + "additionalProperties": true + }, + "yls.snippets.sortMeta": { + "type": "boolean", + "default": true, + "description": "Sort the metadata entries in alphabetical order by keys. Otherwise, insert keys as listed." + }, + "yls.snippets.condition": { + "type": "boolean", + "default": true, + "description": "Enable the condition snippet on YARA rules. Has no effect on the presence of the condition section in the rule snippet." + }, + "yls.snippets.meta": { + "type": "boolean", + "default": true, + "description": "Enable the meta snippet on YARA rules. Has no effect on the presence of the meta section in the rule snippet." + }, + "yls.snippets.rule": { + "type": "boolean", + "default": true, + "description": "Enable the rule skeleton snippet on YARA rules. Settings for other snippets do not affect the sections provided in this snippet." + }, + "yls.snippets.strings": { + "type": "boolean", + "default": true, + "description": "Enable the strings snippet on YARA rules. Has no effect on the presence of the strings section in the rule snippet." } } }, diff --git a/poetry.lock b/poetry.lock index b2427fd..e9a5efc 100644 --- a/poetry.lock +++ b/poetry.lock @@ -480,16 +480,21 @@ toml = ">=0.7.1" [[package]] name = "pytest-yls" -version = "0.1.0" +version = "1.2.1" description = "Pytest plugin to test the YLS as a whole." category = "dev" optional = false -python-versions = ">=3.8,<4.0" +python-versions = "^3.8" +develop = false [package.dependencies] -pytest = ">=7.1.2,<8.0.0" -tenacity = ">=8.0.1,<9.0.0" -yls = ">=1,<2" +pytest = "^7.1.2" +tenacity = "^8.0.1" +yls = "^1" + +[package.source] +type = "directory" +url = "pytest-yls" [[package]] name = "pyyaml" @@ -520,7 +525,7 @@ pbr = ">=2.0.0,<2.1.0 || >2.1.0" [[package]] name = "tenacity" -version = "8.0.1" +version = "8.1.0" description = "Retry code until it succeeds" category = "dev" optional = false @@ -600,7 +605,7 @@ python-versions = ">=3.7" [metadata] lock-version = "1.1" python-versions = ">=3.8,<3.12" -content-hash = "70ec90af03dc9b94202850cd6970b1bdfaf80b33f6dd3ad017a85e5004d9581a" +content-hash = "594dc625822e0fc7037dbc03d83b83b79d321b6876f2b64adebe2858187be113" [metadata.files] astroid = [] diff --git a/pyproject.toml b/pyproject.toml index c1ef6c1..c3c3d1a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,7 +61,7 @@ black = "^22.3.0" isort = "^5.10.1" poethepoet = "^0.13.1" pytest-black = "^0.3.12" -pytest-yls = "^0.1.0" +pytest-yls = {path = "pytest-yls"} [build-system] requires = ["poetry-core>=1.0.0"] @@ -75,6 +75,7 @@ asyncio_mode = "auto" disallow_untyped_calls = false disallow_untyped_defs = true disallow_untyped_decorators = false +disallow_any_generics = false ignore_missing_imports = true strict = true warn_unused_ignores = true @@ -83,6 +84,7 @@ warn_unused_ignores = true good-names = "ls,logger,x,y,c,e,i,j,n,m,f" ignored-classes = "_HookRelay" extension-pkg-allow-list = "yaramod" +ignore-patterns = "^tests/.*" [tool.pylint.messages_control] disable = """, diff --git a/pytest-yls/pytest_yls/plugin.py b/pytest-yls/pytest_yls/plugin.py index 9c3a20c..8300f5b 100644 --- a/pytest-yls/pytest_yls/plugin.py +++ b/pytest-yls/pytest_yls/plugin.py @@ -13,10 +13,14 @@ from pathlib import Path from threading import Thread from typing import Any +from typing import List +from unittest.mock import patch +import pygls.protocol import pytest import pytest_yls.utils as _utils import yaramod +from pygls.lsp import LSP_METHODS_MAP from pygls.lsp import methods from pygls.lsp import types from pygls.server import LanguageServer @@ -61,6 +65,9 @@ class Context: server: YaraLanguageServer cursor_pos: types.Position | None = None + # Config from the client (editor) + config: dict[str, Any] + NOTIFICATION_TIMEOUT_SECONDS = 2.00 CALL_TIMEOUT_SECONDS = 2.00 LANGUAGE_ID = "yara" @@ -71,23 +78,25 @@ def __init__( client: LanguageServer, server: YaraLanguageServer, tmp_path: Path, - config: Any, + pytest_config: Any, files: dict[str, str] | None = None, is_valid_yara_rules_repo: bool = False, + config: dict[str, Any] | None = None, ): self.client = client self.server = server self.cursor_pos = None self.files = files or {} + self.config = config or {} self.is_valid_yara_rules_repo = is_valid_yara_rules_repo self.notification_timeout_seconds = ( - int(config.getoption("yls_notification_timeout")) - if config.getoption("yls_notification_timeout") + int(pytest_config.getoption("yls_notification_timeout")) + if pytest_config.getoption("yls_notification_timeout") else self.NOTIFICATION_TIMEOUT_SECONDS ) self.call_timeout_seconds = ( - int(config.getoption("yls_call_timeout")) - if config.getoption("yls_call_timeout") + int(pytest_config.getoption("yls_call_timeout")) + if pytest_config.getoption("yls_call_timeout") else self.CALL_TIMEOUT_SECONDS ) @@ -107,7 +116,7 @@ def __init__( new_name = yara_rules_root / pathlib.PurePath(name) new_files[str(new_name)] = contents - schema_json = config.stash.get(SCHEMA_JSON_KEY, "") + schema_json = pytest_config.stash.get(SCHEMA_JSON_KEY, "") new_files[str(yara_rules_root / "schema.json")] = schema_json self.files = new_files @@ -138,6 +147,9 @@ def __init__( name = next(iter(self.files)) self.open_file(name) + # Setup the editor configuration + self.client.editor_config = self.config + def open_file(self, name: str) -> None: path = self.tmp_path / name if not path.exists(): @@ -256,9 +268,13 @@ def yls_prepare_with_settings(client_server: Any, tmp_path: Any, pytestconfig) - client, server = client_server def prep( - files: dict[str, str] | None = None, is_valid_yara_rules_repo: bool = False + files: dict[str, str] | None = None, + is_valid_yara_rules_repo: bool = False, + config: dict[str, Any] | None = None, ) -> Context: - return Context(client, server, tmp_path, pytestconfig, files, is_valid_yara_rules_repo) + return Context( + client, server, tmp_path, pytestconfig, files, is_valid_yara_rules_repo, config=config + ) return prep @@ -280,6 +296,22 @@ def hook(ls: LanguageServer, params: Any) -> None: ls.yls_notifications[feature_name].append(params) +def configuration_hook(ls, params): + """WORKSPACE_CONFIGURATION hook""" + editor_config = ls.editor_config + items = params.items + assert len(items) >= 1, "we currently only support single requests" + config = editor_config + item = items[0].section + try: + for part in item.split("."): + config = config[part] + except KeyError: + config = None + + return [config] + + @pytest.fixture(scope="session") def client_server() -> Any: """A fixture to setup a client/server""" @@ -306,8 +338,13 @@ def client_server() -> Any: reset_hooks(client) + client.editor_config = {} + + # Hook configuration requests + client.feature(methods.WORKSPACE_CONFIGURATION)(configuration_hook) + client_thread = Thread( - target=client.start_io, args=(os.fdopen(s2c_r, "rb"), os.fdopen(c2s_w, "wb")) + target=start_editor, args=(client, os.fdopen(s2c_r, "rb"), os.fdopen(c2s_w, "wb")) ) client_thread.daemon = True @@ -317,5 +354,40 @@ def client_server() -> Any: client.lsp.notify(methods.EXIT) server.lsp.notify(methods.EXIT) - server_thread.join() client_thread.join() + server_thread.join() + + +def start_editor(client, stdin, stdout): + """Hook client editor) methods for configuration. + + We need to do this kind of setup because it is really hard to change LSP_METHODS_MAP. + If you want to change the configuration just call Context.set_configuration(). + """ + log.info("[TESTS] Setting up the client (editor) hooks...") + original_deserialize_params = pygls.protocol.deserialize_params + + def _deserialize_params(data, get_params_type): + method = data.get("method") + params = data.get("params") + if method == methods.WORKSPACE_CONFIGURATION and params is not None: + log.warning("[TESTS] We are altering the return value for deserialize_params") + data["params"] = pygls.protocol.dict_to_object(**params) + return data + + return original_deserialize_params(data, get_params_type) + + original_get_method_return_type = pygls.lsp.get_method_return_type + + # pylint: disable=dangerous-default-value + def _get_method_return_type(method_name, lsp_methods_map=LSP_METHODS_MAP): + if method_name == methods.WORKSPACE_CONFIGURATION: + log.warning("[TESTS] We are altering the return value for get_method_return_type") + return List[Any] + + return original_get_method_return_type(method_name, lsp_methods_map) + + patch("pygls.protocol.deserialize_params", _deserialize_params).start() + patch("pygls.protocol.get_method_return_type", _get_method_return_type).start() + + client.start_io(stdin, stdout) diff --git a/tests/unit/test_snippet_string.py b/tests/unit/test_snippet_string.py new file mode 100644 index 0000000..9e4b15c --- /dev/null +++ b/tests/unit/test_snippet_string.py @@ -0,0 +1,27 @@ +# type: ignore + +from yls.snippet_string import SnippetString + + +def test_simple(): + snip = SnippetString("test") + assert str(snip) == "test" + + +def test_placeholder(): + snip = SnippetString() + snip.append_placeholder("one") + snip.append_placeholder("two") + assert str(snip) == "${1:one}${2:two}" + + +def test_complex(): + snip = SnippetString() + snip.append_text("text\n") + snip.append_placeholder("one") + snip.append_tabstop() + snip.append_placeholder("two") + snip.append_variable("CLIPBOARD", "") + snip.append_variable("INVALID_VARIABLE", "default") + snip.append_choice(("c", "ch", "choice")) + assert str(snip) == "text\n${1:one}$2${3:two}${CLIPBOARD}${default}${4|c,ch,choice|}" diff --git a/tests/unit/test_snippets.py b/tests/unit/test_snippets.py new file mode 100644 index 0000000..358e81b --- /dev/null +++ b/tests/unit/test_snippets.py @@ -0,0 +1,89 @@ +# type: ignore + +import pytest + +from yls import snippets +from yls.snippet_string import SnippetString + + +@pytest.mark.parametrize( + "config, expected", + ( + ( + {}, + """rule ${1:my_rule} { +\tmeta: +\t\t${2:KEY} = ${3:"VALUE"} +\tstrings: +\t\t${4:\\$name} = ${5|"string",/regex/,{ HEX }|} +\tcondition: +\t\t${6:any of them} +} +""", + ), + ( + {"metaEntries": {"author": "test user", "hash": ""}}, + """rule ${1:my_rule} { +\tmeta: +\t\tauthor = "test user" +\t\thash = "$2" +\tstrings: +\t\t${3:\\$name} = ${4|"string",/regex/,{ HEX }|} +\tcondition: +\t\t${5:any of them} +} +""", + ), + ( + {"metaEntries": {"filename": "${TM_FILENAME}"}}, + """rule ${1:my_rule} { +\tmeta: +\t\tfilename = "${TM_FILENAME}" +\tstrings: +\t\t${2:\\$name} = ${3|"string",/regex/,{ HEX }|} +\tcondition: +\t\t${4:any of them} +} +""", + ), + ( + { + "metaEntries": { + "author": "", + "date": "${CURRENT_YEAR}-${CURRENT_MONTH}-${CURRENT_DATE}", + } + }, + """rule ${1:my_rule} { +\tmeta: +\t\tauthor = "$2" +\t\tdate = "${CURRENT_YEAR}-${CURRENT_MONTH}-${CURRENT_DATE}" +\tstrings: +\t\t${3:\\$name} = ${4|"string",/regex/,{ HEX }|} +\tcondition: +\t\t${5:any of them} +} +""", + ), + ( + { + "metaEntries": { + "date": "${CURRENT_YEAR}-${CURRENT_MONTH}-${CURRENT_DATE}", + "author": "", + }, + "sortMeta": True, + }, + """rule ${1:my_rule} { +\tmeta: +\t\tauthor = "$2" +\t\tdate = "${CURRENT_YEAR}-${CURRENT_MONTH}-${CURRENT_DATE}" +\tstrings: +\t\t${3:\\$name} = ${4|"string",/regex/,{ HEX }|} +\tcondition: +\t\t${5:any of them} +} +""", + ), + ), +) +def test_basic(config, expected): + assert expected == str(snippets._generate_rule_snippet(SnippetString(), config)) diff --git a/yls/completer.py b/yls/completer.py index 72019f2..bb2f1d3 100644 --- a/yls/completer.py +++ b/yls/completer.py @@ -7,6 +7,7 @@ import pygls.lsp.types as lsp_types from pygls.workspace import Document +from yls import snippets from yls import completion from yls import utils from yls.completion import CONDITION_KEYWORDS @@ -26,8 +27,9 @@ 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)) + async def complete(self, params: lsp_types.CompletionParams) -> lsp_types.CompletionList: + items = await self._complete(params) + return lsp_types.CompletionList(is_incomplete=False, items=items) def signature_help(self, params: lsp_types.CompletionParams) -> lsp_types.SignatureHelp | None: signatures = self._signature_help(params) @@ -71,7 +73,7 @@ def _signature_help( return info - def _complete(self, params: lsp_types.CompletionParams) -> list[lsp_types.CompletionItem]: + async def _complete(self, params: lsp_types.CompletionParams) -> list[lsp_types.CompletionItem]: document = self.ls.workspace.get_document(params.text_document.uri) res = [] @@ -100,6 +102,11 @@ def _complete(self, params: lsp_types.CompletionParams) -> list[lsp_types.Comple log.debug("[COMPLETION] Adding last valid yara file") res += self.complete_last_valid_yara_file(document, params, word) + # Dynamic snippets + log.debug("[COMPLETION] Adding dynamic snippets") + res += await self.complete_dynamic_snippets() + log.debug("[COMPLETION] Adding dynamic snippets done") + # Plugin completion log.debug("COMPLETION] Adding completion items from plugings") res += utils.flatten_list( @@ -238,3 +245,12 @@ def complete_condition_keywords( res.append(item) return res + + async def complete_dynamic_snippets(self) -> list[lsp_types.CompletionItem]: + # Get the configuration each time, in case it changed during the runtime + config = await utils.get_config_from_editor(self.ls, "yls.snippets") + log.debug(f"[COMPLETION] LSP configuration {config=}") + if config is None or not isinstance(config, dict): + return [] + + return snippets.generate_snippets_from_configuration(config) diff --git a/yls/server.py b/yls/server.py index 6351a66..d0abf51 100644 --- a/yls/server.py +++ b/yls/server.py @@ -30,11 +30,11 @@ from yls import icons from yls import linting from yls import utils -from yls.version import __version__ from yls.completer import Completer from yls.hookspecs import ErrorMessage from yls.hover import Hoverer from yls.plugin_manager_provider import PluginManagerProvider +from yls.version import __version__ from yls.yaramod_provider import YaramodProvider log = logging.getLogger(__name__) @@ -227,13 +227,13 @@ def initiliazed(ls: YaraLanguageServer, _params: Any) -> None: @SERVER.feature(COMPLETION, lsp_types.CompletionOptions(trigger_characters=["."])) -def completion( +async def completion( ls: YaraLanguageServer, params: lsp_types.CompletionParams ) -> lsp_types.CompletionList: """Code completion.""" utils.log_command(COMPLETION) - return ls.completer.complete(params) + return await ls.completer.complete(params) @SERVER.feature(SIGNATURE_HELP, lsp_types.SignatureHelpOptions(trigger_characters=["("])) diff --git a/yls/snippet_string.py b/yls/snippet_string.py new file mode 100644 index 0000000..62ba9cb --- /dev/null +++ b/yls/snippet_string.py @@ -0,0 +1,75 @@ +from __future__ import annotations + +from typing import Iterable + +SUPPORTED_VARIABLES = { + "TM_SELECTED_TEXT", + "TM_CURRENT_LINE", + "TM_CURRENT_WORD", + "TM_LINE_INDEX", + "TM_LINE_NUMBER", + "TM_FILENAME", + "TM_FILENAME_BASE", + "TM_DIRECTORY", + "TM_FILEPATH", + "RELATIVE_FILEPATH", + "CLIPBOARD", + "WORKSPACE_NAME", + "WORKSPACE_FOLDER", + "CURSOR_INDEX", + "CURSOR_NUMBER", + "CURRENT_YEAR", + "CURRENT_YEAR_SHORT", + "CURRENT_MONTH", + "CURRENT_MONTH_NAME", + "CURRENT_MONTH_NAME_SHORT", + "CURRENT_DATE", + "CURRENT_DAY_NAME", + "CURRENT_DAY_NAME_SHORT", + "CURRENT_HOUR", + "CURRENT_MINUTE", + "CURRENT_SECOND", + "CURRENT_SECONDS_UNIX", + "RANDOM", + "RANDOM_HEX", + "UUID", + "BLOCK_COMMENT_START", + "BLOCK_COMMENT_END", + "LINE_COMMENT", +} + + +class SnippetString: + cur_idx: int + value: str + + def __init__(self, value: str = ""): + self.value = value + self.cur_idx = 1 + + def append_choice(self, values: Iterable[str]) -> None: + self.value += f"${{{self.get_and_inc()}|{','.join(values)}|}}" + + def append_placeholder(self, value: str) -> None: + self.value += f"${{{self.get_and_inc()}:{value}}}" + + def append_tabstop(self) -> None: + self.value += f"${self.get_and_inc()}" + + def append_text(self, value: str) -> None: + """WARNING: For now you are expected to escape the string if necessary.""" + self.value += value + + def append_variable(self, name: str, default_value: str) -> None: + if name in SUPPORTED_VARIABLES: + self.value += f"${{{name}}}" + else: + self.value += f"${{{default_value}}}" + + def get_and_inc(self) -> int: + i = self.cur_idx + self.cur_idx += 1 + return i + + def __str__(self) -> str: + return self.value diff --git a/yls/snippets.py b/yls/snippets.py new file mode 100644 index 0000000..9e37082 --- /dev/null +++ b/yls/snippets.py @@ -0,0 +1,122 @@ +from __future__ import annotations + +import pygls.lsp.types as lsp_types + +from yls import utils +from yls.snippet_string import SnippetString + + +def generate_snippets_from_configuration( + config: dict | None = None, +) -> list[lsp_types.CompletionItem]: + """Generate dynamic snippets from the user configuration. + + User can specify the configuration in the editor and this function will generate a custom snippet for it. + """ + + # Default is empty config + config = config or {} + + res = [] + if config.get("meta", False): + item = lsp_types.CompletionItem( + label="meta", + detail="Generate a meta section (YARA)", + kind=lsp_types.CompletionItemKind.Snippet, + insert_text=str(_generate_meta_snippet(SnippetString(), config)), + insert_text_format=lsp_types.InsertTextFormat.Snippet, + documentation=utils.markdown_content('meta:\n\tKEY = "VALUE"'), + sort_text="a", + ) + res.append(item) + + if config.get("condition", False): + item = lsp_types.CompletionItem( + label="condition", + detail="Generate a condition section (YARA)", + kind=lsp_types.CompletionItemKind.Snippet, + insert_text=str(_generate_condition_snippet(SnippetString())), + insert_text_format=lsp_types.InsertTextFormat.Snippet, + documentation=utils.markdown_content("condition:\n\tCONDITIONS"), + sort_text="a", + ) + res.append(item) + + if config.get("strings", False): + item = lsp_types.CompletionItem( + label="strings", + detail="Generate a strings skeleton (YARA)", + kind=lsp_types.CompletionItemKind.Snippet, + insert_text=str(_generate_string_snippet(SnippetString())), + insert_text_format=lsp_types.InsertTextFormat.Snippet, + documentation=utils.markdown_content('strings:\n\t$NAME = "STRING"'), + sort_text="a", + ) + res.append(item) + + if config.get("rule", False): + item = lsp_types.CompletionItem( + label="rule", + detail="Generate a rule skeleton (YARA)", + kind=lsp_types.CompletionItemKind.Snippet, + insert_text=str(_generate_rule_snippet(SnippetString(), config)), + insert_text_format=lsp_types.InsertTextFormat.Snippet, + documentation=utils.markdown_content("rule NAME {"), + sort_text="a", + ) + res.append(item) + + return res + + +def _generate_rule_snippet(snippet: SnippetString, config: dict) -> SnippetString: + snippet.append_text("rule ") + snippet.append_placeholder("my_rule") + snippet.append_text(" {\n") + snippet = _generate_meta_snippet(snippet, config) + snippet = _generate_string_snippet(snippet) + snippet.append_text("\n") + snippet = _generate_condition_snippet(snippet) + snippet.append_text("\n}\n") + return snippet + + +def _generate_condition_snippet(snippet: SnippetString) -> SnippetString: + snippet.append_text("\tcondition:\n\t\t") + snippet.append_placeholder("any of them") + return snippet + + +def _generate_string_snippet(snippet: SnippetString) -> SnippetString: + snippet.append_text("\tstrings:\n\t\t") + snippet.append_placeholder(r"\$name") + snippet.append_text(" = ") + snippet.append_choice(('"string"', "/regex/", "{ HEX }")) + return snippet + + +def _generate_meta_snippet(snippet: SnippetString, config: dict) -> SnippetString: + meta_config_dict = config.get("metaEntries", {}) + should_sort_meta = config.get("sortMeta", False) + meta_config: list[tuple[str, str]] = list(meta_config_dict.items()) + if should_sort_meta: + meta_config = sorted(meta_config) + + snippet.append_text("\tmeta:\n") + + if len(meta_config) == 0: + snippet.append_text("\t\t") + snippet.append_placeholder("KEY") + snippet.append_text(" = ") + snippet.append_placeholder('"VALUE"') + snippet.append_text("\n") + else: + for key, value in meta_config: + if value == "": + snippet.append_text(f'\t\t{key} = "') + snippet.append_tabstop() + snippet.append_text('"\n') + else: + snippet.append_text(f'\t\t{key} = "{value}"\n') + + return snippet diff --git a/yls/utils.py b/yls/utils.py index ede9291..1f96f68 100644 --- a/yls/utils.py +++ b/yls/utils.py @@ -82,6 +82,7 @@ def log_command(command: str) -> None: def logging_prolog(plugin_manager: PluginManager) -> None: """Emit dependency version information.""" log.info("Starting yls language server...") + log.info(f"Current file: {__file__}") log.info(f"System platform: {platform.system()}") log.info(f"YLS version: {__version__}") log.info(f"Yaramod version: {yaramod.YARAMOD_VERSION}") @@ -773,7 +774,7 @@ async def get_config_from_editor(ls: Any, section: str) -> Any: ) # We should get list of configuration items here with exactly one element - assert isinstance(config, list) - assert len(config) == 1 + assert isinstance(config, list), "get configuration result should be a list" + assert len(config) >= 1, "did not receive correct configuration response" return config[0]