Skip to content

Commit 01c9a4c

Browse files
ch0mlerch0mler
committed
Merge branch 'variable_definitions'
2 parents e46ace6 + dfd6413 commit 01c9a4c

File tree

4 files changed

+93
-22
lines changed

4 files changed

+93
-22
lines changed

‎.vscode/tasks.json

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
{
22
// See https://go.microsoft.com/fwlink/?LinkId=733558
33
// for the documentation about the tasks.json format
4-
"version": "0.1.0",
4+
"version": "0.1.1",
55
"echoCommand": true,
66
"showOutput": "always",
77
"tasks": [

‎test/extension.test.ts

+29-2
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,9 @@ suite("YARA: Provider", function () {
1616
const filepath: string = path.join(workspace, "peek_rules.yara");
1717
vscode.workspace.openTextDocument(filepath).then(function (doc) {
1818
let defProvider: vscode.DefinitionProvider = new yara.YaraDefinitionProvider();
19-
// SyntaxExample: Line 40, Col 14
19+
// SyntaxExample: Line 42, Col 14
2020
// line numbers start at 0, so we have to subtract one for the lookup
21-
let pos: vscode.Position = new vscode.Position(39, 14);
21+
let pos: vscode.Position = new vscode.Position(41, 14);
2222
let tokenSource: vscode.CancellationTokenSource = new vscode.CancellationTokenSource();
2323
let result = defProvider.provideDefinition(doc, pos, tokenSource.token);
2424
if (result instanceof vscode.Location) {
@@ -41,6 +41,33 @@ suite("YARA: Provider", function () {
4141
});
4242
});
4343

44+
test("variable definition", function (done) {
45+
const filepath: string = path.join(workspace, "peek_rules.yara");
46+
vscode.workspace.openTextDocument(filepath).then(function (doc) {
47+
let defProvider: vscode.DefinitionProvider = new yara.YaraDefinitionProvider();
48+
// $hex_string: Line 25, Col 14
49+
// line numbers start at 0, so we have to subtract one for the lookup
50+
let pos: vscode.Position = new vscode.Position(24, 14);
51+
let tokenSource: vscode.CancellationTokenSource = new vscode.CancellationTokenSource();
52+
let result = defProvider.provideDefinition(doc, pos, tokenSource.token);
53+
if (result instanceof vscode.Location) {
54+
let resultWordRange: vscode.Range = doc.getWordRangeAtPosition(result.range.start);
55+
let resultWord: string = doc.getText(resultWordRange);
56+
if (resultWord == "hex_string") { done(); }
57+
}
58+
else if (result instanceof Array) {
59+
// Should only get one result, so we've failed if an Array is returned
60+
}
61+
else if (result instanceof Promise) {
62+
result.then(function (definition) {
63+
let resultWordRange: vscode.Range = doc.getWordRangeAtPosition(definition.range.start);
64+
let resultWord: string = doc.getText(resultWordRange);
65+
if (resultWord == "hex_string") { done(); }
66+
});
67+
}
68+
});
69+
});
70+
4471
test("rule references", function (done) {
4572
const filepath: string = path.join(workspace, "peek_rules.yara");
4673
vscode.workspace.openTextDocument(filepath).then(function (doc) {

‎test/rules/peek_rules.yara

+3-1
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,8 @@ rule RuleReferenceExample
3636
description = "Rule Reference Example"
3737
author = "Test"
3838
reference = "https://infosec-intern.github.io"
39+
strings:
40+
$hex_string = "test"
3941
condition:
40-
SyntaxExample and "test"
42+
SyntaxExample and $hex_string
4143
}

‎yara/src/extension.ts

+60-18
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,34 @@
33
import * as vscode from "vscode";
44

55

6+
// variables have a few possible first characters - use these to identify vars vs. rules
7+
const varFirstChar = new Set(["$", "#", "@", "!"]);
8+
9+
/*
10+
Get the start and end boundaries for the current YARA rule based on a symbol's position
11+
*/
12+
function GetRuleRange(lines: string[], symbol: vscode.Position) {
13+
let begin: vscode.Position = null;
14+
let end: vscode.Position = null;
15+
const startRuleRegexp = RegExp("^rule ");
16+
const endRuleRegexp = RegExp("^\}");
17+
// only go up to the symbol's line because the rule must be defined before the symbol
18+
for (let lineNo = 0; lineNo < symbol.line; lineNo++) {
19+
if (startRuleRegexp.test(lines[lineNo])) {
20+
begin = new vscode.Position(lineNo, 0);
21+
}
22+
}
23+
// start up this loop again using the beginning of the rule
24+
// and find the line with just a curly brace to identify the end of a rule
25+
for (let lineNo = begin.line; lineNo < lines.length; lineNo++) {
26+
if (endRuleRegexp.test(lines[lineNo])) {
27+
end = new vscode.Position(lineNo, 0);
28+
break;
29+
}
30+
}
31+
return new vscode.Range(begin, end);
32+
}
33+
634
export class YaraDefinitionProvider implements vscode.DefinitionProvider {
735
public provideDefinition(doc: vscode.TextDocument, pos: vscode.Position, token: vscode.CancellationToken): Thenable<vscode.Location> {
836
return new Promise((resolve, reject) => {
@@ -11,21 +39,36 @@ export class YaraDefinitionProvider implements vscode.DefinitionProvider {
1139
const range: vscode.Range = doc.getWordRangeAtPosition(pos);
1240
const symbol: string = doc.getText(range);
1341
// console.log(`Providing definition for symbol '${symbol}'`);
14-
let lines: string[] = doc.getText().split("\n");
15-
let lineNo = 0;
16-
lines.forEach(line => {
17-
let character: number = line.indexOf(symbol);
18-
// line numbers are offset by one in output - need to adjust
19-
// only supporting definitions for rules
20-
if (character != -1 && (lineNo+1) != pos.line && line.startsWith("rule")) {
21-
// console.log(`Found ${symbol} on line ${lineNo} at character ${character}`);
22-
let defPosition: vscode.Position = new vscode.Position(lineNo, character);
23-
definition = new vscode.Location(fileUri, defPosition);
24-
// Definition found. Break out of forEach
25-
return;
42+
let possibleVarStart: vscode.Position = new vscode.Position(range.start.line, range.start.character - 1);
43+
let possibleVarRange: vscode.Range = new vscode.Range(possibleVarStart, range.end);
44+
let possibleVar: string = doc.getText(possibleVarRange);
45+
const lines: string[] = doc.getText().split("\n");
46+
if (varFirstChar.has(possibleVar.charAt(0))) {
47+
// console.log(`Variable detected: ${possibleVar}`);
48+
let currentRuleRange: vscode.Range = GetRuleRange(lines, pos);
49+
// console.log(`Curr rule range: ${currentRuleRange.start.line+1} -> ${currentRuleRange.end.line+1}`);
50+
for (let lineNo = currentRuleRange.start.line; lineNo < currentRuleRange.end.line; lineNo++) {
51+
let character: number = lines[lineNo].indexOf(`$${symbol} =`);
52+
if (character != -1) {
53+
// console.log(`Found defintion of '${possibleVar}' on line ${lineNo + 1} at character ${character + 1}`);
54+
// gotta add one because VSCode won't recognize the '$' as part of the symbol
55+
let defPosition: vscode.Position = new vscode.Position(lineNo, character + 1);
56+
definition = new vscode.Location(fileUri, defPosition);
57+
break;
58+
}
59+
}
60+
}
61+
else {
62+
for (let lineNo = 0; lineNo < pos.line; lineNo++) {
63+
let character: number = lines[lineNo].indexOf(symbol);
64+
if (character != -1 && lines[lineNo].startsWith("rule")) {
65+
// console.log(`Found ${symbol} on line ${lineNo} at character ${character}`);
66+
let defPosition: vscode.Position = new vscode.Position(lineNo, character);
67+
definition = new vscode.Location(fileUri, defPosition);
68+
break;
69+
}
2670
}
27-
lineNo++;
28-
});
71+
}
2972
if (definition != null) {
3073
resolve(definition);
3174
}
@@ -39,14 +82,13 @@ export class YaraDefinitionProvider implements vscode.DefinitionProvider {
3982
export class YaraReferenceProvider implements vscode.ReferenceProvider {
4083
public provideReferences(doc: vscode.TextDocument, pos: vscode.Position, options: { includeDeclaration: boolean }, token: vscode.CancellationToken): Thenable<vscode.Location[]> {
4184
return new Promise((resolve, reject) => {
42-
const varFirstChar = new Set(["$", "#", "@", "!"]);
4385
const fileUri: vscode.Uri = vscode.Uri.file(doc.fileName);
4486
const range: vscode.Range = doc.getWordRangeAtPosition(pos);
4587
let lines: string[] = doc.getText().split("\n");
4688
let references: vscode.Location[] = new Array<vscode.Location>();
4789
let symbol: string = doc.getText(range);
4890
// console.log(`Providing references for symbol '${symbol}'`);
49-
let possibleVarStart: vscode.Position = new vscode.Position(range.start.line, range.start.character-1);
91+
let possibleVarStart: vscode.Position = new vscode.Position(range.start.line, range.start.character - 1);
5092
let possibleVarRange: vscode.Range = new vscode.Range(possibleVarStart, range.end);
5193
let possibleVar: string = doc.getText(possibleVarRange);
5294
if (varFirstChar.has(possibleVar.charAt(0))) {
@@ -57,7 +99,7 @@ export class YaraReferenceProvider implements vscode.ReferenceProvider {
5799
if (character != -1) {
58100
// console.log(`Found ${symbol} on line ${lineNo} at character ${character}`);
59101
// have to readjust the character index
60-
let refPosition: vscode.Position = new vscode.Position(lineNo, character+1);
102+
let refPosition: vscode.Position = new vscode.Position(lineNo, character + 1);
61103
references.push(new vscode.Location(fileUri, refPosition));
62104
}
63105
lineNo++;
@@ -87,7 +129,7 @@ export class YaraReferenceProvider implements vscode.ReferenceProvider {
87129

88130
export function activate(context: vscode.ExtensionContext) {
89131
// console.log("Activating Yara extension");
90-
let YARA: vscode.DocumentSelector = {language: "yara", scheme: "file"};
132+
let YARA: vscode.DocumentSelector = { language: "yara", scheme: "file" };
91133
let definitionDisposable: vscode.Disposable = vscode.languages.registerDefinitionProvider(YARA, new YaraDefinitionProvider());
92134
let referenceDisposable: vscode.Disposable = vscode.languages.registerReferenceProvider(YARA, new YaraReferenceProvider());
93135
context.subscriptions.push(definitionDisposable);

0 commit comments

Comments
 (0)