fixed private rule definitions

added issue #32 test
added isrulestart and isruleend helpers

Author: infosec-intern

test/extension.test.ts

@@ -227,4 +227,34 @@ suite("YARA: Provider", function () {
             }
         });
     });
+
+    /*
+        Trying to capture definitions for private rules
+        Should align symbol `my_private_rule` with definition started by `private rule my_private_rule`
+    */
+   test("issue #32", function (done) {
+        const filepath: string = path.join(workspace, "private_rule_goto.yara");
+        vscode.workspace.openTextDocument(filepath).then(function (doc) {
+            const defProvider: vscode.DefinitionProvider = new YaraDefinitionProvider();
+            // my_private_rule: Line 10, Col 14
+            let pos: vscode.Position = new vscode.Position(9, 14);
+            let tokenSource: vscode.CancellationTokenSource = new vscode.CancellationTokenSource();
+            let result = defProvider.provideDefinition(doc, pos, tokenSource.token);
+            if (result instanceof vscode.Location) {
+                let resultWordRange: vscode.Range = doc.getWordRangeAtPosition(result.range.start);
+                let resultWord: string = doc.getText(resultWordRange);
+                if (resultWord == "my_private_rule" && resultWordRange.start.line == 0 && resultWordRange.start.character == 13) { done(); }
+            }
+            else if (result instanceof Array) {
+                // Should only get one result, so we've failed if an Array is returned
+            }
+            else if (result instanceof Promise) {
+                result.then(function (definition) {
+                    let resultWordRange: vscode.Range = doc.getWordRangeAtPosition(definition.range.start);
+                    let resultWord: string = doc.getText(resultWordRange);
+                    if (resultWord == "my_private_rule" && resultWordRange.start.line == 0 && resultWordRange.start.character == 13) { done(); }
+                });
+            }
+        });
+   });
 });

test/rules/private_rule_goto.yara

@@ -0,0 +1,11 @@
+private rule my_private_rule
+{
+    condition:
+        true
+}
+
+rule my_public_rule
+{
+    condition:
+        my_private_rule
+}

yara/src/definitionProvider.ts

@@ -1,7 +1,7 @@
 "use strict";
 
 import * as vscode from "vscode";
-import {GetRuleRange, varFirstChar} from "./helpers";
+import {GetRuleRange, varFirstChar, IsRuleStart} from "./helpers";
 
 
 export class YaraDefinitionProvider implements vscode.DefinitionProvider {
@@ -34,7 +34,7 @@ export class YaraDefinitionProvider implements vscode.DefinitionProvider {
             else {
                 for (let lineNo = 0; lineNo < pos.line; lineNo++) {
                     let character: number = lines[lineNo].indexOf(symbol);
-                    if (character != -1 && lines[lineNo].startsWith("rule")) {
+                    if (character != -1 && IsRuleStart(lines[lineNo])) {
                         // console.log(`Found ${symbol} on line ${lineNo} at character ${character}`);
                         let defPosition: vscode.Position = new vscode.Position(lineNo, character);
                         definition = new vscode.Location(fileUri, defPosition);

yara/src/helpers.ts

@@ -12,22 +12,34 @@ export const varFirstChar: Set<string> = new Set(["$", "#", "@", "!"]);
 export function GetRuleRange(lines: string[], symbol: vscode.Position) {
     let begin: vscode.Position | null = null;
     let end: vscode.Position | null = null;
-    const startRuleRegexp = RegExp("^rule ");
-    const endRuleRegexp = RegExp("^\}");
     // find the nearest reference to "rule" by traversing the lines in reverse order
     for (let lineNo = symbol.line; lineNo >= 0; lineNo--) {
-        if (startRuleRegexp.test(lines[lineNo])) {
+        if (IsRuleStart(lines[lineNo])) {
             begin = new vscode.Position(lineNo, 0);
             break;
         }
     }
     // start up this loop again using the beginning of the rule
     // and find the line with just a curly brace to identify the end of a rule
     for (let lineNo = begin.line; lineNo < lines.length; lineNo++) {
-        if (endRuleRegexp.test(lines[lineNo])) {
+        if (IsRuleEnd(lines[lineNo])) {
             end = new vscode.Position(lineNo, 0);
             break;
         }
     }
     return new vscode.Range(begin, end);
-}
+}
+
+/*
+    Determine if the given line is the start of a YARA rule or not
+*/
+export function IsRuleStart(line: string) {
+    return new RegExp("^(private )?rule ").test(line);
+}
+
+/*
+    Determine if the given line is the end of a YARA rule or not
+*/
+export function IsRuleEnd(line: string) {
+    return new RegExp("^\}").test(line);
+}