Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Prompt Security's extension integration #784

Closed
wants to merge 17 commits into from
5 changes: 4 additions & 1 deletion package.json
Original file line number Diff line number Diff line change
Expand Up @@ -872,8 +872,8 @@
"@typescript-eslint/eslint-plugin": "^7.0.1",
"@typescript-eslint/parser": "^7.2.0",
"chai": "4.3.1",
"eslint-config-prettier": "^9.1.0",
"eslint": "^8.57.0",
"eslint-config-prettier": "^9.1.0",
"mocha": "10.3.0",
"typescript": "^5.4.2",
"vsce": "^2.15.0",
Expand All @@ -885,9 +885,12 @@
"dependencies": {
"@checkmarxdev/ast-cli-javascript-wrapper": "0.0.86",
"copyfiles": "2.4.1",
"esbuild": "^0.20.1",
"eslint-config-prettier": "^9.1.0",
"eslint-plugin-node": "^11.1.0",
"eslintcc": "^0.8.1",
"express": "^4.18.3",
"http": "^0.0.1-security",
"tree-kill": "^1.2.2"
},
"overrides": {
Expand Down
18 changes: 17 additions & 1 deletion src/extension.ts
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,11 @@ import { WorkspaceListener } from "./utils/listener/workspaceListener";
import { DocAndFeedbackView } from "./views/docsAndFeedbackView/docAndFeedbackView";
import { messages } from "./utils/common/messages";
import { commands } from "./utils/common/commands";
import { PromptSecurity } from "./utils/listener/promptSecurity";



let promptListener: PromptSecurity | null = null;

export async function activate(context: vscode.ExtensionContext) {
// Create logs channel and make it visible
Expand Down Expand Up @@ -203,8 +208,19 @@ export async function activate(context: vscode.ExtensionContext) {
kicsScanCommand.registerKicsRemediation();
// Refresh sca tree with start scan message
scaResultsProvider.refreshData(constants.scaStartScan);

//You will need to add a checkbox to enable the extension's activation
const isPromptEnabled:boolean = true;
if (isPromptEnabled){
//The port number should be dynamic
promptListener = new PromptSecurity(context,3312);
}
}

// eslint-disable-next-line @typescript-eslint/no-empty-function
export function deactivate() {
}
//Close Prompt's server if the extension is deactivated
if (promptListener){
promptListener.deactivate();
}
}
64 changes: 64 additions & 0 deletions src/test/8.promptSecurity.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,64 @@
import * as vscode from 'vscode';
import * as assert from 'assert';
import { PromptSecurity } from '../utils/listener/promptSecurity';

suite('Extension Test Suite', () => {
vscode.window.showInformationMessage('Start all tests.');
const mockContext: vscode.ExtensionContext = {
globalState: {
get: () => undefined,
update: () => Promise.resolve()
},
// Other properties of ExtensionContext are omitted for brevity
} as any;
let promptSecurity = new PromptSecurity(mockContext, 3000);
test('getSelectedText returns the selected text', async () => {
// Set up a mock text editor with a selection
const document = await vscode.workspace.openTextDocument({ content: 'Hello, world!' });
const editor = await vscode.window.showTextDocument(document);
const selection = new vscode.Selection(0, 0, 0, 5);
editor.selection = selection;

// Simulate a selection change event
const event: vscode.TextEditorSelectionChangeEvent = {
textEditor: editor,
selections: [selection],
kind: undefined
};

// Test the function
const selectedText = promptSecurity.getSelectedText(event);
assert.strictEqual(selectedText, 'Hello');
});

test('getLastSelectedCodeSections returns an empty array initially', () => {
// Set up a mock context


// Bind the mock context to the function
const boundGetLastSelectedCodeSections = promptSecurity.getLastSelectedCodeSections.bind({ context: mockContext });

// Test the function
const codeSections = boundGetLastSelectedCodeSections();
assert.deepStrictEqual(codeSections, []);
});
test('getLastSelectedCodeSections returns an array with the last selection', () => {
// Set up a mock context
let lastSelectedCodeSections = [];
lastSelectedCodeSections.push({
code: "code",
filePath: "filePath"
});
mockContext.globalState.update('lastSelectedCodeSections', lastSelectedCodeSections);

// Bind the mock context to the function
const boundGetLastSelectedCodeSections = promptSecurity.getLastSelectedCodeSections.bind({ context: mockContext });

// Test the function
const codeSections = boundGetLastSelectedCodeSections();
assert.deepStrictEqual(codeSections, [{
code: "code",
filePath: "filePath"
}]);
});
});
123 changes: 123 additions & 0 deletions src/utils/listener/promptSecurity.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,123 @@
import * as vscode from 'vscode';
//For Prompt's server
import express from 'express';
import { Server } from "http";

const MIN_CODE_LENGTH = 20;
const MAX_CODE_LENGTH = 32 * 1024;

export class PromptSecurity {
private hostname: string;
private port?: number;
private server?: Server;
private app?: express;
private context?:vscode.ExtensionContext;

constructor(context:vscode.ExtensionContext,port:number) {
this.hostname = "127.0.0.1";
this.context = context;
this.port = port;
this.registerPromptListener();
}
getServer(){
return this.server;
}
deactivate(){
if (this.server){
this.server.close();
}
}
registerPromptListener(){
try {
guy-ps marked this conversation as resolved.
Show resolved Hide resolved
this.app = express();
this.app.use(express.json());
this.app.use(express.urlencoded({ extended: true }));
this.server = this.app.listen(this.port,this.hostname);
//Starting the endpoint the browser extension can call
this.extensionListener();
//On selection send the context and the event to the selection buffer funtion
vscode.window.onDidChangeTextEditorSelection(event => {
this.selectionBuffer(event);
});
//If the window is changed - Prompt
vscode.window.onDidChangeWindowState(async windowState => {
if (windowState){
this.windowChange();
}
});
} catch(error) {
vscode.window.showErrorMessage(`An error occurred: ${error.message}`);
}
}
extensionListener(
){
guy-ps marked this conversation as resolved.
Show resolved Hide resolved
this.app.get('/api/health', (req,res) => {
res.json({ "vscode":"Yay!" });
});
this.app.post('/api/checkCode', (req, res) => {
const receivedText: string = req.body.text ?? "";
let containsCode: boolean = false;
let filePath: string | undefined = undefined;
const lastSelectedCodeSections: SelectionBuffer[] | undefined = this.context.globalState.get('lastSelectedCodeSections');
if (lastSelectedCodeSections) {
for (const section of lastSelectedCodeSections) {
if (receivedText.includes(section.code)) {
containsCode = true;
filePath = section.filePath;
break;
}
}
}
res.json({ containsCode, filePath });
});
}
windowChange(){
if (!this.server || !this.server.address()) {
this.server = this.app.listen(this.port, this.hostname);
}
}
getSelectedText(
event:vscode.TextEditorSelectionChangeEvent
){
let code = "";
if (event.selections.length > 0) {
const selection = event.selections[0];
if (!selection.isEmpty) {
code = event.textEditor.document.getText(selection)?.trim();
if (code.length < MIN_CODE_LENGTH || code.length > MAX_CODE_LENGTH) {
code = "";
}
}
}
return code;
}
getLastSelectedCodeSections(){
let lastSelectedCodeSections: SelectionBuffer[] | undefined = this.context.globalState.get('lastSelectedCodeSections');
if (!lastSelectedCodeSections) {
lastSelectedCodeSections = [];
}
if (lastSelectedCodeSections.length > 25) {
lastSelectedCodeSections.shift();
}
return lastSelectedCodeSections;
}
selectionBuffer(
event:vscode.TextEditorSelectionChangeEvent
){
const code = this.getSelectedText(event);
if (code !==""){
const filePath = event.textEditor.document.uri.fsPath ?? "";
const lastSelectedCodeSections = this.getLastSelectedCodeSections();
lastSelectedCodeSections.push({
code: code,
filePath: filePath
});
this.context.globalState.update('lastSelectedCodeSections', lastSelectedCodeSections);
}
}
}
interface SelectionBuffer {
code: string;
filePath:string;
}

Loading