Dynamic snippets

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."
         }
       }
     },

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 = """,

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)

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|}"

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))