Merge pull request #5344 from ethereum/pastedCodeSafety

Pasted code safety

apps/remix-ide/src/app/files/fileManager.ts

@@ -26,7 +26,7 @@ const profile = {
     'readFile', 'copyFile', 'copyDir', 'rename', 'mkdir', 'readdir', 'dirList', 'fileList', 'remove', 'getCurrentFile', 'getFile',
     'getFolder', 'setFile', 'switchFile', 'refresh', 'getProviderOf', 'getProviderByName', 'getPathFromUrl', 'getUrlFromPath',
     'saveCurrentFile', 'setBatchFiles', 'isGitRepo', 'isFile', 'isDirectory', 'hasGitSubmodule', 'copyFolderToJson', 'diff',
-    'hasGitSubmodules'
+    'hasGitSubmodules', 'getOpenedFiles'
   ],
   kind: 'file-system'
 }

apps/remix-ide/src/app/plugins/remixAIPlugin.tsx

@@ -3,7 +3,7 @@ import { ViewPlugin } from '@remixproject/engine-web'
 import { Plugin } from '@remixproject/engine';
 import { RemixAITab, ChatApi } from '@remix-ui/remix-ai'
 import React, { useCallback } from 'react';
-import { ICompletions, IModel, RemoteInferencer, IRemoteModel, IParams, GenerationParams, CodeExplainAgent } from '@remix/remix-ai-core';
+import { ICompletions, IModel, RemoteInferencer, IRemoteModel, IParams, GenerationParams, CodeExplainAgent, SecurityAgent } from '@remix/remix-ai-core';
 import { CustomRemixApi } from '@remix-api'
 import { PluginViewWrapper } from '@remix-ui/helper'
 const _paq = (window._paq = window._paq || [])
@@ -17,9 +17,8 @@ const profile = {
   displayName: 'RemixAI',
   methods: ['code_generation', 'code_completion',
     "solidity_answer", "code_explaining",
-    "code_insertion", "error_explaining",
-    "initialize", 'chatPipe', 'ProcessChatRequestBuffer',
-    'isChatRequestPending'],
+    "code_insertion", "error_explaining", "vulnerability_check",
+    "initialize", 'chatPipe', 'ProcessChatRequestBuffer', 'isChatRequestPending'],
   events: [],
   icon: 'assets/img/remix-logo-blue.png',
   description: 'RemixAI provides AI services to Remix IDE.',
@@ -38,15 +37,16 @@ export class RemixAIPlugin extends ViewPlugin {
   remoteInferencer:RemoteInferencer = null
   isInferencing: boolean = false
   chatRequestBuffer: chatRequestBufferT<any> = null
-  agent: CodeExplainAgent
+  codeExpAgent: CodeExplainAgent
+  securityAgent: SecurityAgent
   useRemoteInferencer:boolean = false
   dispatch: any
 
   constructor(inDesktop:boolean) {
     super(profile)
     this.isOnDesktop = inDesktop
-    this.agent = new CodeExplainAgent(this)
-    // user machine dont use resource for remote inferencing
+    this.codeExpAgent = new CodeExplainAgent(this)
+    // user machine dont use ressource for remote inferencing
   }
 
   onActivation(): void {
@@ -62,6 +62,8 @@ export class RemixAIPlugin extends ViewPlugin {
       this.useRemoteInferencer = true
       this.initialize()
     }
+
+    this.securityAgent = new SecurityAgent(this)
   }
 
   async initialize(model1?:IModel, model2?:IModel, remoteModel?:IRemoteModel, useRemote?:boolean){
@@ -97,11 +99,6 @@ export class RemixAIPlugin extends ViewPlugin {
   }
 
   async code_generation(prompt: string): Promise<any> {
-    if (this.isInferencing) {
-      this.call('terminal', 'log', { type: 'aitypewriterwarning', value: "RemixAI is already busy!" })
-      return
-    }
-
     if (this.isOnDesktop && !this.useRemoteInferencer) {
       return await this.call(this.remixDesktopPluginName, 'code_generation', prompt)
     } else {
@@ -118,17 +115,8 @@ export class RemixAIPlugin extends ViewPlugin {
   }
 
   async solidity_answer(prompt: string, params: IParams=GenerationParams): Promise<any> {
-    if (this.isInferencing) {
-      this.call('terminal', 'log', { type: 'aitypewriterwarning', value: "RemixAI is already busy!" })
-      return
-    }
-    if (prompt.trimStart().startsWith('gpt') || prompt.trimStart().startsWith('sol-gpt')) {
-      params.terminal_output = true
-      params.stream_result = false
-      params.return_stream_response = false
-    }
+    const newPrompt = await this.codeExpAgent.chatCommand(prompt)
 
-    const newPrompt = await this.agent.chatCommand(prompt)
     let result
     if (this.isOnDesktop && !this.useRemoteInferencer) {
       result = await this.call(this.remixDesktopPluginName, 'solidity_answer', newPrompt)
@@ -142,11 +130,6 @@ export class RemixAIPlugin extends ViewPlugin {
   }
 
   async code_explaining(prompt: string, context: string, params: IParams=GenerationParams): Promise<any> {
-    if (this.isInferencing) {
-      this.call('terminal', 'log', { type: 'aitypewriterwarning', value: "RemixAI is already busy!" })
-      return
-    }
-
     let result
     if (this.isOnDesktop && !this.useRemoteInferencer) {
       result = await this.call(this.remixDesktopPluginName, 'code_explaining', prompt, context, params)
@@ -159,11 +142,6 @@ export class RemixAIPlugin extends ViewPlugin {
   }
 
   async error_explaining(prompt: string, context: string="", params: IParams=GenerationParams): Promise<any> {
-    if (this.isInferencing) {
-      this.call('terminal', 'log', { type: 'aitypewriterwarning', value: "RemixAI is already busy!" })
-      return
-    }
-
     let result
     if (this.isOnDesktop && !this.useRemoteInferencer) {
       result = await this.call(this.remixDesktopPluginName, 'error_explaining', prompt)
@@ -174,6 +152,22 @@ export class RemixAIPlugin extends ViewPlugin {
     return result
   }
 
+  async vulnerability_check(prompt: string, params: IParams=GenerationParams): Promise<any> {
+    let result
+    if (this.isOnDesktop && !this.useRemoteInferencer) {
+      result = await this.call(this.remixDesktopPluginName, 'vulnerability_check', prompt)
+
+    } else {
+      result = await this.remoteInferencer.vulnerability_check(prompt, params)
+    }
+    if (result && params.terminal_output) this.call('terminal', 'log', { type: 'aitypewriterwarning', value: result })
+    return result
+  }
+
+  getVulnerabilityReport(file: string): any {
+    return this.securityAgent.getReport(file)
+  }
+
   async code_insertion(msg_pfx: string, msg_sfx: string): Promise<any> {
     if (this.isOnDesktop && !this.useRemoteInferencer) {
       return await this.call(this.remixDesktopPluginName, 'code_insertion', msg_pfx, msg_sfx)
@@ -194,11 +188,12 @@ export class RemixAIPlugin extends ViewPlugin {
         if (fn === "code_explaining") ChatApi.composer.send("Explain the current code")
         else if (fn === "error_explaining") ChatApi.composer.send("Explain the error")
         else if (fn === "solidity_answer") ChatApi.composer.send("Answer the following question")
-        else console.log("chatRequestBuffer is not empty. First process the last request.")
+        else if (fn === "vulnerability_check") ChatApi.composer.send("Is there any vulnerability in the pasted code?")
+        else console.log("chatRequestBuffer function name not recognized.")
       }
     }
     else {
-      console.log("chatRequestBuffer is not empty. First process the last request.")
+      console.log("chatRequestBuffer is not empty. First process the last request.", this.chatRequestBuffer)
     }
     _paq.push(['trackEvent', 'ai', 'remixAI_chat', 'askFromTerminal'])
   }

apps/remix-ide/src/app/tabs/locales/en/editor.json

@@ -28,6 +28,7 @@
   "editor.explainFunctionByAI": "```\n{content}\n```\nExplain the function {currentFunction}",
   "editor.explainFunctionByAISol": "```\n{content}\n```\nExplain the function {currentFunction}",
   "editor.ExplainPipeMessage": "```\n {content}\n```\nExplain the snipped above",
+  "editor.PastedCodeSafety": "```\n {content}\n```\n\nReply in a short manner: Does this code contain major security vulnerabilities leading to a scam or loss of funds?",
   "editor.executeFreeFunction": "Run a free function",
   "editor.executeFreeFunction2": "Run the free function \"{name}\"",
   "editor.toastText1": "This can only execute free function",

libs/remix-ai-core/src/agents/securityAgent.ts

@@ -1,28 +1,188 @@
 // security checks
 import * as fs from 'fs';
 
-class SecurityAgent {
-  private codebase: string[]; // list of codebase files
+interface SecurityReport {
+  compiled: boolean;
+  vulnerabilities: string[];
+  isNotSafe: string;
+  fileName: string;
+  reportTimestamp: string;
+  recommendations: string[];
+  fileModifiedSinceLastReport: boolean;
+  hasPastedCode: boolean;
+}
+
+class WorkspaceWatcher {
+  private intervalId: NodeJS.Timeout | null = null;
+  public interval: number;
+  private task: () => void;
+
+  constructor(task: () => void, interval: number) {
+    this.task = task;
+    this.interval = interval;
+  }
+
+  start(): void {
+    if (this.intervalId === null) {
+      this.intervalId = setInterval(() => {
+        this.task();
+      }, this.interval);
+    }
+  }
+
+  stop(): void {
+    if (this.intervalId !== null) {
+      clearInterval(this.intervalId);
+      this.intervalId = null;
+    }
+  }
+
+  isRunning(): boolean {
+    return this.intervalId !== null;
+  }
+}
+
+export class SecurityAgent {
   public currentFile: string;
+  public openedFiles: any;
+  private basePlugin: any;
+  private watcher: WorkspaceWatcher;
+  public reports: SecurityReport[] = [];
+
+  constructor(plugin) {
+    this.basePlugin = plugin;
+
+    this.basePlugin.on('fileManager', 'fileAdded', (path) => { });
+    this.basePlugin.on('fileManager', 'fileChanged', (path) => { //this.modifiedFile(path)
+    });
+
+    this.basePlugin.on('fileManager', 'fileRemoved', (path) => { this.removeFileFromReport(path) });
+    this.basePlugin.on('fileManager', 'fileRenamed', (oldName, newName) => {
+      this.removeFileFromReport(oldName);
+      this.addFileToReport(newName);
+    });
+
+    this.basePlugin.on('solidity', 'compilationFinished', async (fileName, source, languageVersion, data) => { this.onCompilationFinished(fileName) });
+    this.basePlugin.on('vyper', 'compilationFinished', async (fileName, source, languageVersion, data) => { this.onCompilationFinished(fileName) });
+    this.basePlugin.on('hardhat', 'compilationFinished', async (fileName, source, languageVersion, data) => { this.onCompilationFinished(fileName) });
+    this.basePlugin.on('foundry', 'compilationFinished', async (fileName, source, languageVersion, data) => { this.onCompilationFinished(fileName) });
+
+    this.watcher = new WorkspaceWatcher(async () => {
+      try {
+        this.currentFile = await this.basePlugin.call('fileManager', 'getCurrentFile');
+        this.openedFiles = await this.basePlugin.call('fileManager', 'getOpenedFiles');
+
+        Object.keys(this.openedFiles).forEach(key => {
+          this.addFileToReport(this.openedFiles[key]);
+        });
+      } catch (error) {
+        // no file selected or opened currently
+      }
+    }, 10000);
+    this.watcher.start();
+  }
+
+  addFileToReport(file: string): void {
+    const report = this.reports.find((r) => r.fileName === file);
+    if (report) {
+      // nothing to do
+    } else {
+      this.reports.push({
+        compiled: false,
+        isNotSafe: 'No',
+        vulnerabilities: [],
+        fileName: file,
+        reportTimestamp: null,
+        recommendations: [],
+        fileModifiedSinceLastReport: false,
+        hasPastedCode: false
+      });
+    }
 
-  constructor(codebasePath: string) {
-    // git or fs
-    this.codebase = this.loadCodebase(codebasePath);
   }
 
-  private loadCodebase(path: string): string[] {
-    const files = fs.readdirSync(path);
-    return files
-      .filter(file => file.endsWith('.ts'))
-      .flatMap(file => fs.readFileSync(`${path}/${file}`, 'utf-8').split('\n'));
+  async onCompilationFinished(file: string) {
+    let report = this.reports.find((r) => r.fileName === file);
+    if (report) {
+      report.compiled = true;
+      report.fileModifiedSinceLastReport = false;
+    } else {
+      report = {
+        compiled: true,
+        isNotSafe: 'No',
+        vulnerabilities: [],
+        fileName: file,
+        reportTimestamp: null,
+        recommendations: [],
+        fileModifiedSinceLastReport: false,
+        hasPastedCode: false
+      }
+      this.reports.push(report);
+    }
+
+    try {
+      this.processFile(file);
+      console.log('Checking for vulnerabilities after compilation', this.reports);
+    } catch (error) {
+      console.error('Error checking for vulnerabilities after compilation: ', error);
+    }
+
+    // check for security vulnerabilities
   }
 
-  public update(currentFile, lineNumber){
+  removeFileFromReport(file: string): void {
+    const index = this.reports.findIndex((r) => r.fileName === file);
+    if (index !== -1) {
+      this.reports.splice(index, 1);
+    }
+  }
+
+  modifiedFile(file: string): void {
+    const report = this.reports.find((r) => r.fileName === file);
+    if (report) {
+      report.fileModifiedSinceLastReport = true;
+    }
+  }
+
+  async processFile(file: string) {
+    try {
+      const report = this.reports.find((r) => r.fileName === file);
+      if (report) { }
+      else {
+        this.reports.push({
+          compiled: false,
+          isNotSafe: 'No',
+          vulnerabilities: [],
+          fileName: file,
+          reportTimestamp: null,
+          recommendations: [],
+          fileModifiedSinceLastReport: false,
+          hasPastedCode: false
+        });
+      }
+
+      if (!report.reportTimestamp || report.fileModifiedSinceLastReport) {
+        const content = await this.basePlugin.call('fileManager', 'getFile', file);
+        const prompt = "```\n" + content + "\n```\n\nReply in a short manner: Does this code contain major security vulnerabilities leading to a scam or loss of funds?"
+
+        let result = await this.basePlugin.call('remixAI', 'vulnerability_check', prompt)
+        result = JSON.parse(result);
+        report.vulnerabilities = result.Reason;
+        report.recommendations = result.Suggestion;
+        report.isNotSafe = result.Answer;
+        report.reportTimestamp = new Date().toISOString();
+      }
+
+    } catch (error) {
+      console.error('Error processing file: ', error);
+    }
+  }
 
+  getReport(file: string): SecurityReport {
+    return this.reports.find((r) => r.fileName === file);
   }
 
   public getRecommendations(currentLine: string, numSuggestions: number = 3): string[] {
-    // process the codebase highlighting security vulnerabilities and deliver recommendations
     const suggestions: string[] = [];
     return suggestions;
   }

libs/remix-ai-core/src/helpers/streamHandler.ts

@@ -46,6 +46,7 @@ export const HandleStreamResponse = async (streamResponse,
       }
       catch (error) {
         console.error('Error parsing JSON:', error);
+        return { 'generateText': 'Try again!', 'isGenerating': false }
       }
     }
     if (done_cb) {
@@ -54,7 +55,7 @@ export const HandleStreamResponse = async (streamResponse,
   }
   catch (error) {
     console.error('Error parsing JSON:', error);
-    return { 'generateText': '', 'isGenerating': false }
+    return { 'generateText': 'Try again!', 'isGenerating': false }
   }
 }
 

libs/remix-ai-core/src/index.ts

@@ -21,4 +21,5 @@ export {
 
 export * from './types/types'
 export * from './helpers/streamHandler'
-export * from './agents/codeExplainAgent'
+export * from './agents/codeExplainAgent'
+export * from './agents/securityAgent'

libs/remix-api/src/lib/plugins/remixAIDesktop-api.ts

@@ -10,7 +10,7 @@ export interface IRemixAID {
 
   } & StatusEvents,
   methods: {
-    code_completion(context: string): Promise<string> 
+    code_completion(context: string): Promise<string>
     code_insertion(msg_pfx: string, msg_sfx: string): Promise<string>,
     code_generation(prompt: string): Promise<string | null>,
     code_explaining(code: string, context?: string): Promise<string | null>,

libs/remix-ui/editor/src/lib/providers/inlineCompletionProvider.ts

@@ -9,7 +9,7 @@ export class RemixInLineCompletionProvider implements monacoTypes.languages.Inli
   props: EditorUIProps
   monaco: any
   completionEnabled: boolean
-  task: string
+  task: string = 'code_completion'
   currentCompletion: any
   private lastRequestTime: number = 0;
   private readonly minRequestInterval: number = 200;
@@ -60,13 +60,6 @@ export class RemixInLineCompletionProvider implements monacoTypes.languages.Inli
       endColumn: getTextAtLine(model.getLineCount()).length + 1,
     });
 
-    if (!word.endsWith(' ') &&
-      !word.endsWith('.') &&
-      !word.endsWith('"') &&
-      !word.endsWith('(')) {
-      return;
-    }
-
     try {
       const split = word.split('\n')
       if (split.length < 2) return

libs/remix-ui/editor/src/lib/remix-ui-editor.tsx

@@ -2,7 +2,7 @@ import React, { useState, useRef, useEffect, useReducer } from 'react' // eslint
 import { FormattedMessage, useIntl } from 'react-intl'
 import { isArray } from 'lodash'
 import Editor, { DiffEditor, loader, Monaco } from '@monaco-editor/react'
-import { AlertModal } from '@remix-ui/app'
+import { AppModal } from '@remix-ui/app'
 import { ConsoleLogs, QueryParams } from '@remix-project/remix-lib'
 import { reducerActions, reducerListener, initialState } from './actions/editor'
 import { solidityTokensProvider, solidityLanguageConfig } from './syntaxes/solidity'
@@ -664,11 +664,26 @@ export const EditorUI = (props: EditorUIProps) => {
       }
     })
 
-    editor.onDidPaste((e) => {
+    editor.onDidPaste(async (e) => {
       if (!pasteCodeRef.current && e && e.range && e.range.startLineNumber >= 0 && e.range.endLineNumber >= 0 && e.range.endLineNumber - e.range.startLineNumber > 10) {
-        const modalContent: AlertModal = {
+        // get the file name
+        const pastedCode = editor.getModel().getValueInRange(e.range)
+        const pastedCodePrompt = intl.formatMessage({ id: 'editor.PastedCodeSafety' }, { content:pastedCode })
+
+        const modalContent: AppModal = {
           id: 'newCodePasted',
-          title: intl.formatMessage({ id: 'editor.title1' }),
+          title: "New code pasted",
+          okLabel: 'Ask RemixAI',
+          cancelLabel: 'Close',
+          cancelFn: () => {},
+          okFn: async () => {
+            await props.plugin.call('popupPanel', 'showPopupPanel', true)
+            setTimeout(async () => {
+              props.plugin.call('remixAI', 'chatPipe', 'vulnerability_check', pastedCodePrompt)
+            }, 500)
+            // add matamo event
+            _paq.push(['trackEvent', 'ai', 'remixAI', 'vulnerability_check_pasted_code'])
+          },
           message: (
             <div>
               {' '}
@@ -699,10 +714,9 @@ export const EditorUI = (props: EditorUIProps) => {
                 </div>
               </div>
             </div>
-          ),
+          )
         }
-        props.plugin.call('notification', 'alert', modalContent)
-        pasteCodeRef.current = true
+        props.plugin.call('notification', 'modal', modalContent)
         _paq.push(['trackEvent', 'editor', 'onDidPaste', 'more_than_10_lines'])
       }
     })