Skip to content

Commit 955bbba

Browse files
ch0mlerch0mler
committed
moved providers out of extension.ts
getrulerange and varFirstChar => helpers DefinitionProvider => definitionProvider ReferenceProvider => referenceProvider
1 parent 3c11195 commit 955bbba

File tree

5 files changed

+167
-149
lines changed

5 files changed

+167
-149
lines changed

‎test/extension.test.ts

+8-6
Original file line numberDiff line numberDiff line change
@@ -7,16 +7,18 @@ Please refer to their documentation on https://mochajs.org/ for help.
77

88
import * as path from "path";
99
import * as vscode from "vscode";
10-
import * as yara from "../yara/src/extension";
1110
import { YaraCompletionItemProvider } from "../yara/src/completionProvider";
11+
import { YaraDefinitionProvider } from "../yara/src/definitionProvider";
12+
import { YaraReferenceProvider } from "../yara/src/referenceProvider";
13+
1214

1315
let workspace = path.join(__dirname, "..", "..", "test/rules/");
1416

1517
suite("YARA: Provider", function () {
1618
test("rule definition", function (done) {
1719
const filepath: string = path.join(workspace, "peek_rules.yara");
1820
vscode.workspace.openTextDocument(filepath).then(function (doc) {
19-
const defProvider: vscode.DefinitionProvider = new yara.YaraDefinitionProvider();
21+
const defProvider: vscode.DefinitionProvider = new YaraDefinitionProvider();
2022
// SyntaxExample: Line 43, Col 14
2123
// line numbers start at 0, so we have to subtract one for the lookup
2224
let pos: vscode.Position = new vscode.Position(42, 14);
@@ -45,7 +47,7 @@ suite("YARA: Provider", function () {
4547
test("variable definition", function (done) {
4648
const filepath: string = path.join(workspace, "peek_rules.yara");
4749
vscode.workspace.openTextDocument(filepath).then(function (doc) {
48-
const defProvider: vscode.DefinitionProvider = new yara.YaraDefinitionProvider();
50+
const defProvider: vscode.DefinitionProvider = new YaraDefinitionProvider();
4951
// $hex_string: Line 25, Col 14
5052
// line numbers start at 0, so we have to subtract one for the lookup
5153
let pos: vscode.Position = new vscode.Position(24, 14);
@@ -72,7 +74,7 @@ suite("YARA: Provider", function () {
7274
test("symbol references", function (done) {
7375
const filepath: string = path.join(workspace, "peek_rules.yara");
7476
vscode.workspace.openTextDocument(filepath).then(function (doc) {
75-
const refProvider: vscode.ReferenceProvider = new yara.YaraReferenceProvider();
77+
const refProvider: vscode.ReferenceProvider = new YaraReferenceProvider();
7678
// $dstring: Line 22, Col 11
7779
let pos: vscode.Position = new vscode.Position(21, 11);
7880
// console.log(`search term: ${doc.getText(doc.getWordRangeAtPosition(pos))}`);
@@ -110,7 +112,7 @@ suite("YARA: Provider", function () {
110112
test("wildcard references", function (done) {
111113
const filepath: string = path.join(workspace, "peek_rules.yara");
112114
vscode.workspace.openTextDocument(filepath).then(function (doc) {
113-
const refProvider: vscode.ReferenceProvider = new yara.YaraReferenceProvider();
115+
const refProvider: vscode.ReferenceProvider = new YaraReferenceProvider();
114116
// $hex_*: Line 31, Col 11
115117
let pos: vscode.Position = new vscode.Position(30, 11);
116118
// console.log(`search term: ${doc.getText(doc.getWordRangeAtPosition(pos))}`);
@@ -193,7 +195,7 @@ suite("YARA: Provider", function () {
193195
test("issue #17", function (done) {
194196
const filepath: string = path.join(workspace, "peek_rules.yara");
195197
vscode.workspace.openTextDocument(filepath).then(function (doc) {
196-
let refProvider: vscode.ReferenceProvider = new yara.YaraReferenceProvider();
198+
let refProvider: vscode.ReferenceProvider = new YaraReferenceProvider();
197199
// $hex_string: Line 20, Col 11
198200
let pos: vscode.Position = new vscode.Position(19, 11);
199201
// console.log(`search term: ${doc.getText(doc.getWordRangeAtPosition(pos))}`);

‎yara/src/definitionProvider.ts

+53
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
"use strict";
2+
3+
import * as vscode from "vscode";
4+
import {GetRuleRange, varFirstChar} from "./helpers";
5+
6+
7+
export class YaraDefinitionProvider implements vscode.DefinitionProvider {
8+
public provideDefinition(doc: vscode.TextDocument, pos: vscode.Position, token: vscode.CancellationToken): Thenable<vscode.Location> {
9+
return new Promise((resolve, reject) => {
10+
let definition: vscode.Location | null = null;
11+
const fileUri: vscode.Uri = vscode.Uri.file(doc.fileName);
12+
const range: vscode.Range = doc.getWordRangeAtPosition(pos);
13+
const symbol: string = doc.getText(range);
14+
// console.log(`Providing definition for symbol '${symbol}'`);
15+
let possibleVarStart: vscode.Position = new vscode.Position(range.start.line, range.start.character - 1);
16+
let possibleVarRange: vscode.Range = new vscode.Range(possibleVarStart, range.end);
17+
let possibleVar: string = doc.getText(possibleVarRange);
18+
const lines: string[] = doc.getText().split("\n");
19+
if (varFirstChar.has(possibleVar.charAt(0))) {
20+
// console.log(`Variable detected: ${possibleVar}`);
21+
let currentRuleRange: vscode.Range = GetRuleRange(lines, pos);
22+
// console.log(`Curr rule range: ${currentRuleRange.start.line+1} -> ${currentRuleRange.end.line+1}`);
23+
for (let lineNo = currentRuleRange.start.line; lineNo < currentRuleRange.end.line; lineNo++) {
24+
let character: number = lines[lineNo].indexOf(`$${symbol} =`);
25+
if (character != -1) {
26+
// console.log(`Found defintion of '${possibleVar}' on line ${lineNo + 1} at character ${character + 1}`);
27+
// gotta add one because VSCode won't recognize the '$' as part of the symbol
28+
let defPosition: vscode.Position = new vscode.Position(lineNo, character + 1);
29+
definition = new vscode.Location(fileUri, defPosition);
30+
break;
31+
}
32+
}
33+
}
34+
else {
35+
for (let lineNo = 0; lineNo < pos.line; lineNo++) {
36+
let character: number = lines[lineNo].indexOf(symbol);
37+
if (character != -1 && lines[lineNo].startsWith("rule")) {
38+
// console.log(`Found ${symbol} on line ${lineNo} at character ${character}`);
39+
let defPosition: vscode.Position = new vscode.Position(lineNo, character);
40+
definition = new vscode.Location(fileUri, defPosition);
41+
break;
42+
}
43+
}
44+
}
45+
if (definition != null) {
46+
resolve(definition);
47+
}
48+
else {
49+
reject();
50+
}
51+
});
52+
}
53+
}

‎yara/src/extension.ts

+2-143
Original file line numberDiff line numberDiff line change
@@ -2,151 +2,10 @@
22

33
import * as vscode from "vscode";
44
import {YaraCompletionItemProvider} from "./completionProvider";
5+
import {YaraDefinitionProvider} from "./definitionProvider";
6+
import {YaraReferenceProvider} from "./referenceProvider";
57

68

7-
// variables have a few possible first characters - use these to identify vars vs. rules
8-
const varFirstChar: Set<string> = new Set(["$", "#", "@", "!"]);
9-
10-
/*
11-
Get the start and end boundaries for the current YARA rule based on a symbol's position
12-
*/
13-
function GetRuleRange(lines: string[], symbol: vscode.Position) {
14-
let begin: vscode.Position | null = null;
15-
let end: vscode.Position | null = null;
16-
const startRuleRegexp = RegExp("^rule ");
17-
const endRuleRegexp = RegExp("^\}");
18-
// find the nearest reference to "rule" by traversing the lines in reverse order
19-
for (let lineNo = symbol.line; lineNo >= 0; lineNo--) {
20-
if (startRuleRegexp.test(lines[lineNo])) {
21-
begin = new vscode.Position(lineNo, 0);
22-
break;
23-
}
24-
}
25-
// start up this loop again using the beginning of the rule
26-
// and find the line with just a curly brace to identify the end of a rule
27-
for (let lineNo = begin.line; lineNo < lines.length; lineNo++) {
28-
if (endRuleRegexp.test(lines[lineNo])) {
29-
end = new vscode.Position(lineNo, 0);
30-
break;
31-
}
32-
}
33-
return new vscode.Range(begin, end);
34-
}
35-
36-
export class YaraDefinitionProvider implements vscode.DefinitionProvider {
37-
public provideDefinition(doc: vscode.TextDocument, pos: vscode.Position, token: vscode.CancellationToken): Thenable<vscode.Location> {
38-
return new Promise((resolve, reject) => {
39-
let definition: vscode.Location | null = null;
40-
const fileUri: vscode.Uri = vscode.Uri.file(doc.fileName);
41-
const range: vscode.Range = doc.getWordRangeAtPosition(pos);
42-
const symbol: string = doc.getText(range);
43-
// console.log(`Providing definition for symbol '${symbol}'`);
44-
let possibleVarStart: vscode.Position = new vscode.Position(range.start.line, range.start.character - 1);
45-
let possibleVarRange: vscode.Range = new vscode.Range(possibleVarStart, range.end);
46-
let possibleVar: string = doc.getText(possibleVarRange);
47-
const lines: string[] = doc.getText().split("\n");
48-
if (varFirstChar.has(possibleVar.charAt(0))) {
49-
// console.log(`Variable detected: ${possibleVar}`);
50-
let currentRuleRange: vscode.Range = GetRuleRange(lines, pos);
51-
// console.log(`Curr rule range: ${currentRuleRange.start.line+1} -> ${currentRuleRange.end.line+1}`);
52-
for (let lineNo = currentRuleRange.start.line; lineNo < currentRuleRange.end.line; lineNo++) {
53-
let character: number = lines[lineNo].indexOf(`$${symbol} =`);
54-
if (character != -1) {
55-
// console.log(`Found defintion of '${possibleVar}' on line ${lineNo + 1} at character ${character + 1}`);
56-
// gotta add one because VSCode won't recognize the '$' as part of the symbol
57-
let defPosition: vscode.Position = new vscode.Position(lineNo, character + 1);
58-
definition = new vscode.Location(fileUri, defPosition);
59-
break;
60-
}
61-
}
62-
}
63-
else {
64-
for (let lineNo = 0; lineNo < pos.line; lineNo++) {
65-
let character: number = lines[lineNo].indexOf(symbol);
66-
if (character != -1 && lines[lineNo].startsWith("rule")) {
67-
// console.log(`Found ${symbol} on line ${lineNo} at character ${character}`);
68-
let defPosition: vscode.Position = new vscode.Position(lineNo, character);
69-
definition = new vscode.Location(fileUri, defPosition);
70-
break;
71-
}
72-
}
73-
}
74-
if (definition != null) {
75-
resolve(definition);
76-
}
77-
else {
78-
reject();
79-
}
80-
});
81-
}
82-
}
83-
84-
export class YaraReferenceProvider implements vscode.ReferenceProvider {
85-
public provideReferences(doc: vscode.TextDocument, pos: vscode.Position, options: { includeDeclaration: boolean }, token: vscode.CancellationToken): Thenable<vscode.Location[]> {
86-
return new Promise((resolve, reject) => {
87-
const fileUri: vscode.Uri = vscode.Uri.file(doc.fileName);
88-
const range: vscode.Range = doc.getWordRangeAtPosition(pos);
89-
let lines: string[] = doc.getText().split("\n");
90-
let references: vscode.Location[] = new Array<vscode.Location>();
91-
let symbol: string = doc.getText(range);
92-
// console.log(`Providing references for symbol '${symbol}'`);
93-
let possibleVarStart: vscode.Position = new vscode.Position(range.start.line, range.start.character - 1);
94-
let possibleVarRange: vscode.Range = new vscode.Range(possibleVarStart, range.end);
95-
let possibleVar: string = doc.getText(possibleVarRange);
96-
if (varFirstChar.has(possibleVar.charAt(0))) {
97-
let varRegexp: string;
98-
let startLine: number;
99-
let endLine: number;
100-
// console.log(`Identified symbol as a variable: ${symbol}`);
101-
let possibleWildcardEnd: vscode.Position = new vscode.Position(range.end.line, range.end.character + 1);
102-
let possibleWildcardRange: vscode.Range = new vscode.Range(possibleVarStart, possibleWildcardEnd);
103-
let possibleWildcard: string = doc.getText(possibleWildcardRange);
104-
if (possibleWildcard.slice(-1) == "*") {
105-
// treat like a wildcard and search only the local rule
106-
varRegexp = `[\$#@!]${symbol}[a-zA-Z0-9_]+`;
107-
let ruleRange = GetRuleRange(lines, pos);
108-
startLine = ruleRange.start.line;
109-
endLine = ruleRange.end.line;
110-
}
111-
else {
112-
// treat like a normal variable reference and search the whole document
113-
varRegexp = `[\$#@!]${symbol}[^a-zA-Z0-9_]`;
114-
startLine = 0;
115-
endLine = lines.length;
116-
}
117-
for (let lineNo = startLine; lineNo < endLine; lineNo++) {
118-
let character: number = lines[lineNo].search(varRegexp);
119-
if (character != -1) {
120-
// console.log(`Found ${symbol} on line ${lineNo} at character ${character}`);
121-
// have to readjust the character index
122-
let refPosition: vscode.Position = new vscode.Position(lineNo, character + 1);
123-
references.push(new vscode.Location(fileUri, refPosition));
124-
}
125-
126-
}
127-
}
128-
else {
129-
let lineNo = 0;
130-
lines.forEach(line => {
131-
let character: number = line.indexOf(symbol);
132-
if (character != -1) {
133-
// console.log(`Found ${symbol} on line ${lineNo} at character ${character}`);
134-
let refPosition: vscode.Position = new vscode.Position(lineNo, character);
135-
references.push(new vscode.Location(fileUri, refPosition));
136-
}
137-
lineNo++;
138-
});
139-
}
140-
if (references != null) {
141-
resolve(references);
142-
}
143-
else {
144-
reject();
145-
}
146-
});
147-
}
148-
}
149-
1509
export function activate(context: vscode.ExtensionContext) {
15110
// console.log("Activating Yara extension");
15211
let YARA: vscode.DocumentSelector = { language: "yara", scheme: "file" };

‎yara/src/helpers.ts

+33
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
"use strict";
2+
3+
import * as vscode from "vscode";
4+
5+
6+
// variables have a few possible first characters - use these to identify vars vs. rules
7+
export const varFirstChar: Set<string> = new Set(["$", "#", "@", "!"]);
8+
9+
/*
10+
Get the start and end boundaries for the current YARA rule based on a symbol's position
11+
*/
12+
export function GetRuleRange(lines: string[], symbol: vscode.Position) {
13+
let begin: vscode.Position | null = null;
14+
let end: vscode.Position | null = null;
15+
const startRuleRegexp = RegExp("^rule ");
16+
const endRuleRegexp = RegExp("^\}");
17+
// find the nearest reference to "rule" by traversing the lines in reverse order
18+
for (let lineNo = symbol.line; lineNo >= 0; lineNo--) {
19+
if (startRuleRegexp.test(lines[lineNo])) {
20+
begin = new vscode.Position(lineNo, 0);
21+
break;
22+
}
23+
}
24+
// start up this loop again using the beginning of the rule
25+
// and find the line with just a curly brace to identify the end of a rule
26+
for (let lineNo = begin.line; lineNo < lines.length; lineNo++) {
27+
if (endRuleRegexp.test(lines[lineNo])) {
28+
end = new vscode.Position(lineNo, 0);
29+
break;
30+
}
31+
}
32+
return new vscode.Range(begin, end);
33+
}

‎yara/src/referenceProvider.ts

+71
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
"use strict";
2+
3+
import * as vscode from "vscode";
4+
import {GetRuleRange, varFirstChar} from "./helpers";
5+
6+
7+
export class YaraReferenceProvider implements vscode.ReferenceProvider {
8+
public provideReferences(doc: vscode.TextDocument, pos: vscode.Position, options: { includeDeclaration: boolean }, token: vscode.CancellationToken): Thenable<vscode.Location[]> {
9+
return new Promise((resolve, reject) => {
10+
const fileUri: vscode.Uri = vscode.Uri.file(doc.fileName);
11+
const range: vscode.Range = doc.getWordRangeAtPosition(pos);
12+
let lines: string[] = doc.getText().split("\n");
13+
let references: vscode.Location[] = new Array<vscode.Location>();
14+
let symbol: string = doc.getText(range);
15+
// console.log(`Providing references for symbol '${symbol}'`);
16+
let possibleVarStart: vscode.Position = new vscode.Position(range.start.line, range.start.character - 1);
17+
let possibleVarRange: vscode.Range = new vscode.Range(possibleVarStart, range.end);
18+
let possibleVar: string = doc.getText(possibleVarRange);
19+
if (varFirstChar.has(possibleVar.charAt(0))) {
20+
let varRegexp: string;
21+
let startLine: number;
22+
let endLine: number;
23+
// console.log(`Identified symbol as a variable: ${symbol}`);
24+
let possibleWildcardEnd: vscode.Position = new vscode.Position(range.end.line, range.end.character + 1);
25+
let possibleWildcardRange: vscode.Range = new vscode.Range(possibleVarStart, possibleWildcardEnd);
26+
let possibleWildcard: string = doc.getText(possibleWildcardRange);
27+
if (possibleWildcard.slice(-1) == "*") {
28+
// treat like a wildcard and search only the local rule
29+
varRegexp = `[\$#@!]${symbol}[a-zA-Z0-9_]+`;
30+
let ruleRange = GetRuleRange(lines, pos);
31+
startLine = ruleRange.start.line;
32+
endLine = ruleRange.end.line;
33+
}
34+
else {
35+
// treat like a normal variable reference and search the whole document
36+
varRegexp = `[\$#@!]${symbol}[^a-zA-Z0-9_]`;
37+
startLine = 0;
38+
endLine = lines.length;
39+
}
40+
for (let lineNo = startLine; lineNo < endLine; lineNo++) {
41+
let character: number = lines[lineNo].search(varRegexp);
42+
if (character != -1) {
43+
// console.log(`Found ${symbol} on line ${lineNo} at character ${character}`);
44+
// have to readjust the character index
45+
let refPosition: vscode.Position = new vscode.Position(lineNo, character + 1);
46+
references.push(new vscode.Location(fileUri, refPosition));
47+
}
48+
49+
}
50+
}
51+
else {
52+
let lineNo = 0;
53+
lines.forEach(line => {
54+
let character: number = line.indexOf(symbol);
55+
if (character != -1) {
56+
// console.log(`Found ${symbol} on line ${lineNo} at character ${character}`);
57+
let refPosition: vscode.Position = new vscode.Position(lineNo, character);
58+
references.push(new vscode.Location(fileUri, refPosition));
59+
}
60+
lineNo++;
61+
});
62+
}
63+
if (references != null) {
64+
resolve(references);
65+
}
66+
else {
67+
reject();
68+
}
69+
});
70+
}
71+
}

0 commit comments

Comments
 (0)