Prompt Security's extension integration

package.json

@@ -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.3",
     "vsce": "^2.15.0",
@@ -885,9 +885,12 @@
   "dependencies": {
     "@checkmarxdev/ast-cli-javascript-wrapper": "0.0.87",
     "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": {

src/extension.ts

@@ -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
@@ -203,8 +208,22 @@ export async function activate(context: vscode.ExtensionContext) {
     kicsScanCommand.registerKicsRemediation();
     // Refresh sca tree with start scan message
     scaResultsProvider.refreshData(constants.scaStartScan);
+    try{
+        //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);
+        }
+    } catch(error) {
+        logs.error(`An error occurred while registering the promptListener: ${error.message}\n It is not active`);
+    }
 }
 
 // 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();
+    }
+}

src/test/8.promptSecurity.test.ts

@@ -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"
+		}]);
+    });
+});

src/utils/listener/promptSecurity.ts

@@ -0,0 +1,154 @@
+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) {
+		try{
+			this.hostname = "127.0.0.1";
+			this.context = context;
+			this.port = port;
+			this.registerPromptListener();
+		} catch(error) {
+			vscode.window.showErrorMessage(`An error occurred while activating the PromptListener: ${error.message}`);
+		}
+	}
+	getServer(){
+		return this.server;
+	}
+	deactivate(){
+		try{
+			if (this.server){
+				this.server.close();
+			}
+		} catch(error) {
+			vscode.window.showErrorMessage(`An error occurred while deactivating the plugin: ${error.message}`);
+		}
+	}
+	registerPromptListener(){
+		try {
+			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 while registering the promptListener: ${error.message}`);
+		}
+	}
+	extensionListener(
+	){
+		this.app.get('/api/health', (req,res) => {
+			res.json({ "vscode":"Yay!" });
+		});
+		this.app.post('/api/checkCode', (req, res) => {
+			try{
+				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 });
+			} catch(error) {
+				vscode.window.showErrorMessage(`An error occurred while checking if th code was copied: ${error.message}`);
+				res.json({ containsCode: false });
+			}
+		});
+	}
+	windowChange(){
+		try{
+			if (!this.server || !this.server.address()) {
+				this.server = this.app.listen(this.port, this.hostname);
+			}
+		} catch(error) {
+			vscode.window.showErrorMessage(`An error occurred in the process of restarting the server due to window change: ${error.message}`);
+		}
+	}
+	getSelectedText(
+		event:vscode.TextEditorSelectionChangeEvent
+	){
+		try{	
+			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;
+		} catch(error) {
+			vscode.window.showErrorMessage(`An error occurred while getting selected text: ${error.message}`);
+			return "";
+		}
+	}
+	getLastSelectedCodeSections(){
+		try{
+			let lastSelectedCodeSections: SelectionBuffer[] | undefined = this.context.globalState.get('lastSelectedCodeSections');
+			if (!lastSelectedCodeSections) {
+				lastSelectedCodeSections = [];
+			}
+			if (lastSelectedCodeSections.length > 25) {
+				lastSelectedCodeSections.shift();
+			}
+			return lastSelectedCodeSections;
+		} catch(error) {
+			vscode.window.showErrorMessage(`An error occurred while getting last selected code sections: ${error.message}`);
+			return [];
+		}
+	}
+	selectionBuffer(
+		event:vscode.TextEditorSelectionChangeEvent
+	){
+		try{
+			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);
+			}
+		} catch(error) {	
+			vscode.window.showErrorMessage(`An error occurred while updating the selection buffer: ${error.message}`);
+		}
+	}
+}
+interface SelectionBuffer {
+	code: string;
+	filePath:string;
+}
+